Column/row slicing a torch sparse tensor

1

I have a pytorch sparse tensor that I need sliced row/column wise using this slice [idx][:,idx] where idx is a list of indexes, using the mentioned slice yields my desired result on an ordinary float tensor. Is it possible applying the same slicing on a sparse tensor? Example here:

#constructing sparse matrix
i = np.array([[0,1,2,2],[0,1,2,1]])
v = np.ones(4)
i = torch.from_numpy(i.astype("int64"))
v = torch.from_numpy(v.astype("float32"))
test1 = torch.sparse.FloatTensor(i, v)

#constructing float tensor
test2 = np.array([[1,0,0],[0,1,0],[0,1,1]])
test2 = autograd.Variable(torch.cuda.FloatTensor(test2), requires_grad=False)

#slicing
idx = [1,2]
print(test2[idx][:,idx])

output:

Variable containing:
 1  0
 1  1
[torch.cuda.FloatTensor of size 2x2 (GPU 0)]

I am holding a 250.000 x 250.000 adjacency matrix, where I need to slice 5000 rows and 5000 columns, using the random idx, by simply sampling 5000 random idx’s. Since the dataset is so large it is not realistic to convert to a more convenient datatype.

can I achieve the same slicing result on test1? Is it even possible? If not, are there any work-arounds?

2

Pytorch currently does not support slicing of sparse tensors. One thing you could do is look at the sparse tensor’s tensor._indices() and then pick out the correct indices that correspond to the indices you want to slice.

3

Is it possible to use this approach so everything is running on the GPU? right now I am doing as you are saying but I am using np.where() I only get 12% GPU utlization running my code. This most likely happens because I alternate between having my data on CPU and GPU