Numpy "put_along_axis" equivalent in torch

1

Hi,

Is there a pytorch equivalent of numpy’s put_along_axis functionality?
I am trying to make values of a tensor to 1 based on given indices in another tensor. Rightnow, I am doing this using numpy and wondering if there is any better way to do this Pytorch?

2

I think .scatter_ might be a good approach.
If you post example tensors, we could have a look, if that’s a good option.

1 Like
3

Thanks for your reply.
What I wanted is something like this:

zero_tensor = torch.zeros([3,4])
index_tensor = tensor([[1, 2, 0], [0, 0, 1], [0,2,2]])

result = tensor([[1., 1., 1., 0.],
[1., 1., 0., 0.],
[1., 0., 1., 0.]])

replace values of zero_tensor from 0 to 1, where the indices are given by the index_tensor. And this operation has to be performed row-wise.
“scatter_” solves this problem. Now I don’t need to use put_along_axis from numpy. Thanks once again.

1 Like