Assign value to a tensor based on an index tensor?

1

Hi, I’m trying to use index tensor to assign value to another tensor. In particular, I have tensor A of shape [a,b,c] and index tensor B of shape [a,b].

I have a value tensor C of shape [b]
I want to assign values such that A[i,j,B[i,j]] = C[j].

I tried using a for loop but it’s too slow. Is there a proper approach to do this ? (I also checked some tutorial on tensor indexing e.g. here and here but they aren’t relatable)

2

Could you post the (slow) reference implementation, please?

3

Thank for the response. I came up with the code

import torch
a = torch.rand((3,3,3))
b = torch.randint(0,3,(3,3)).unsqueeze(-1)
c = torch.full_like(b,fill_value=5,dtype=a.dtype)
print(a)
print(b)
print(a.scatter_(dim=2,index=b,src=c))

which did what I wanted