Pytorch equivalent for


I’m looking for pytorch equivalent for but couldn’t find any option. I found the “index_add_” similar as 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.arange(2)[:,None,None],np.arange(3)[:,None],index[:,None,:]), to_multiply)


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:

1 Like