Pytorch equivalent for np.multiply.at

Hi,

I’m looking for pytorch equivalent for np.multiply.at but couldn’t find any option. I found the “index_add_” similar as np.add.at but nothing for multiply. Wondering if there is a efficient way to do select an index tensor (may have same index) from a base tensor and then multiply with another tensor.

In Numpy, this is how I achieve it:

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]])
np.multiply.at(base1, (np.arange(2)[:,None,None],np.arange(3)[:,None],index[:,None,:]), to_multiply)

Thanks!

Would this library work? I think spmm supports gradients, but spspmm does not. Might want to double check.

1 Like

@aluo-x thanks! though it doesn’t seem to fit my needs, but when searching for similar packages I found torch_scatter.scatter_mul and it seems to work: https://pytorch-scatter.readthedocs.io/en/1.4.0/functions/mul.html

1 Like