Using tensordot with torch.sparse tensors

Question

Is it possible to use a similar method as "tensordot" with torch.sparse tensors?

I am trying to apply a 4 dimensional tensor onto a 2 dimensional tensor. This is possible using torch or numpy. However, I did not find the way to do it using torch.sparse without making the sparse tensor dense using ".to_dense()".

More precisely, here is what I want to do without using ".to_dense()":

import torch
import torch.sparse

nb_x = 4
nb_y = 3
coordinates = torch.LongTensor([[0,1,2],[0,1,2],[0,1,2],[0,1,2]])
values = torch.FloatTensor([1,2,3])
tensor4D = torch.sparse.FloatTensor(coordinates,values,torch.Size([nb_x,nb_y,nb_x,nb_y]))
inp = torch.rand((nb_x,nb_y))

#what I want to do
out = torch.tensordot(tensor4D.to_dense(),inp,dims=([2,3],[0,1]))

print(inp)
print(out)

(here is the output: torch_code)

Alternatively, here is a similar code using numpy:

import numpy as np

tensor4D = np.zeros((4,3,4,3))
tensor4D[0,0,0,0] = 1
tensor4D[1,1,1,1] = 2
tensor4D[2,2,2,2] = 3
inp = np.random.rand(4,3)

out = np.tensordot(tensor4D,inp)

print(inp)
print(out)

(here is the output: numpy_code)

Thanks for helping!

Answer 1

Your specific tensordot can be cast to a simple matrix multiplication by "squeezing" the first two and last two dimensions of tensor4D.

In short, what you want to do is

raw = tensor4D.view(nb_x*nb_y, nb_x*nb_y) @ inp.flatten()
out = raw.view(nb_x, nb_y)

However, since view and reshape are not implemented for sparse tensors, you'll have to it manually:

sz = tensor4D.shape
coeff = torch.tensor([[1, sz[1], 0, 0], [0, 0, 1, sz[3]]])
reshaped = torch.sparse.FloatTensor(coeff @ idx, tensor4D._values(), torch.Size([nb_x*nb_y, nb_x*nb_y]))

# once we reshaped tensord4D it's all downhill from here
raw = torch.sparse.mm(reshaped, inp.flatten()[:, None])
out = raw.reshape(nb_x, nb_y)
print(out)

And the output is

tensor([[0.4180, 0.0000, 0.0000],
   [0.0000, 0.6025, 0.0000],
   [0.0000, 0.0000, 0.5897],
   [0.0000, 0.0000, 0.0000]])

Answer 2

Indeed, this works very well, thank you for your answer!

The weakness of this method seems to me that it is hard to generalize.

In fact, "inp" and "out" are supposed to be images. Here, they are black and white images since there are only two dimensions: height and width.

If instead, I take RGB images, then I will have to consider 6D tensors acting on 3D tensors. I can still apply the same trick by "squeezing" the first three dimensions together and the last three dimensions together. However it seems to me that it will become more involving very quickly (maybe I am wrong). While using tensordot instead would be much more simpler for generalization.

Therefore, I am going to use the solution you proposed, but I am still interested if someone finds an other solution.

Answer 3

In another answer, a reshaping coeff matrix is proposed for shaping the indices of the 4D matrix:

raw = tensor4D.view(nb_x*nb_y, nb_x*nb_y) @ inp.flatten()
out = raw.view(nb_x, nb_y)
sz = tensor4D.shape
coeff = torch.tensor([[1, sz[1], 0, 0], [0, 0, 1, sz[3]]])
reshaped = torch.sparse.FloatTensor(coeff @ idx, tensor4D._values(), torch.Size([nb_x*nb_y, nb_x*nb_y]))
raw = torch.sparse.mm(reshaped, inp.flatten()[:, None])
out = raw.reshape(nb_x, nb_y)

I was able to get the same answer as tensordot if I change the coeff to be instead:

coeff = torch.tensor([[sz[1], 1, 0, 0], [0, 0, sz[3], 1]])

This is because flattening a 2D (x, y) coordinate to 1D would be x*dim(y) + y rather than x + dim(y)*y.

(This is a comment on Shai's answer but since I'm starting a new account I would have to start an independent answer).