2

Are there equivalent of np.multiply.at in Pytorch? I have two 4d arrays and one 2d index array:

base = torch.ones((2, 3, 5, 5))
to_multiply = torch.arange(120).view(2, 3, 4, 5)
index = torch.tensor([[0, 2, 4, 2], [0, 3, 3, 2]])

As shown in this question I asked earlier (in Numpy), the row index of the index array corresponds to the 1st dimension of base and to_multiply, and the value of the index array corresponds to the 3rd dimension of base. I want to take the slice from base according to the index and multiply with to_multiply, it can be achieved in Numpy as follows:

np.multiply.at(base1, (np.arange(2)[:,None,None],np.arange(3)[:,None],index[:,None,:]), to_multiply)

However, now when I want to translate this to PyTorch, I cannot find an equivalent of np.multiply.at in Pytorch, I can only find the "index_add_" method but there is no "index_multiply". And I want to avoid doing explicit for loop.

So how can I achieve above in PyTorch? Thanks!

MachineLearner
  • 413
  • 5
  • 10
  • 1
    In case anyone else have same question, found the solution: the torch-scatter package has a scatter_mul that does exactly this: https://pytorch-scatter.readthedocs.io/en/1.4.0/functions/mul.html – MachineLearner Jul 31 '20 at 15:54
  • Check this out: https://stackoverflow.com/a/65584479/3337089 – Nagabhushan S N Jan 05 '21 at 18:36
  • depending on the values you actually plan to multiply, and the precision you need, you could always take the log first, then apply index_add, then exponentiate the result. – Erik Jun 19 '21 at 15:53

0 Answers0