Todo functions and autograd supports for Sparse Tensor

[weiyangfb]

Here summarizes a list of requested Sparse Tensor functions and autograd supports from previous PRs. Please feel free to comment on functions that should be added also.

Functions

• sum() with autograd [sparse] torch.sparse.sum() #12430
• max() with autograd
• log1p() Add log1p for sparse tensor #8969
• S.copy_(S) with autograd copy_(Sparse, Sparse) for sparse tensor #9005
• indexing (gather(), index_select())
• mul_(S, D) -> S, mul(S, D) -> S with autograd
• cuda()
• nn.Linear with autograd (SxS, SxD, relies on addmm and matmul)
• softmax() with autograd (same as in TF: Applies softmax() to a region of a densified tensor submatrix; (2) Masks out the zero locations; (3) Renormalizes the remaining elements. SparseTensor result has exactly the same non-zero indices and shape)
• to_sparse() dense.to_sparse() re: #8853 #12171
• narrow_copy() add narrow() support for sparse tensors re: #8853 #11342
• sparse_mm(S, D) -> D with autograd
• cat() implement concatenation of sparse tensors #13577
• unsqueeze(), stack() Implement unsqueeze for sparse vectors (this also makes stack work out of the box) #13760

Wish list

• bmm(S, D) (add an extra sparse dim at indices of SparseTensor as batch dim?)
• broadcasting mul(S, D) -> S
• Dataset, Dataloader
• save, load for sparse tensors

Existing

• autograd supported for values() via [sparse] Autograd indices/values and sparse_coo ctor #13001 (Thanks to @ssnl!), that means all element-wise ops are supported in sparse now
• norm (cannot take dim args)
• pow
• clone
• zero_
• t_ / t
• add_ / add(Sparse, Sparse, Scalar) -> Sparse
• add_ / add(Dense, Sparse, Scalar) -> Dense
• sub_ / sub(Sparse, Sparse, Scalar) -> Sparse
• mul_ / mul(Sparse, Sparse) -> Sparse
• mul_ / mul(Sparse, Scalar) -> Sparse
• div_ / div(Sparse, Scalar) -> Sparse
• addmm(Dense, Sparse, Dense, Scalar, Scalar) -> Dense
• sspaddmm(Sparse, Sparse, Dense, Scalar, Scalar) -> Sparse
• mm(Sparse, Dense) -> Dense
• smm(Sparse, Dense) -> Sparse
• hspmm(Sparse, Dense) -> HybridSparse
• spmm(Sparse, Dense) -> Dense

[ezyang]

Also:

• Feature request: sparse x sparse -> sparse #1550
• [Feature Request] Sparse-Dense elementwise Multiplication  #3158

Also Wojciej has been posting about these recently. He'd like some variant of this script to work:

# sparse matrix (e.g. weights)
indices = torch.LongTensor([[0,2], [1,0], [1,2]])
indices = indices.transpose(1,0).contiguous()
values = torch.FloatTensor(int(indices.view(-1).size(0)/2)).uniform_(-1,1)
t = torch.sparse.FloatTensor(indices, values, torch.Size([10,10])).cuda()

# sparse vector (e.g. activations)
v = torch.sparse.FloatTensor(
    torch.LongTensor([[0], [1], [2]]).transpose(1,0), 
    torch.FloatTensor([3,4,5]),
    torch.Size([10])
).cuda()

# works fine
print(F.linear(v.to_dense(), t.to_dense()))

# RuntimeError: numel is not implemented for type torch.cuda.sparse.FloatTensor
# print(F.linear(v, t.to_dense()))

# RuntimeError: numel is not implemented for type torch.cuda.sparse.FloatTensor
# print(F.linear(v, t)

# RuntimeError: Expected object of type torch.cuda.FloatTensor but found type torch.cuda.sparse.FloatTensor for argument #2 'mat2'
# print(F.linear(v.to_dense(), t))

[ezyang]

See also #8856

[weiyangfb]

@ezyang Ok I am adding them as well. For #1550, is it asking for matrix-matrix multiplication or element-wise multiplication, or both?

[weiyangfb]

Requests added:

• x.copy_(y) doesn't work with sparse x #8330 copy_()
• sparse.mm(S, D) with autograd #2389 mm(S, D) autograd (e.g. sparse_mm(S, D) with sparse gradient)
• torch.sparse.FloatTensor implementation requests #4710 sum(), log1p(), max()
• TypeError: Type torch.sparse.LongTensor doesn't implement stateless method cat #3248 cat(S, S)
• [feature request] select_index for sparse tensors #7416 indexing (gather, index_select)
• [feature request] sparse x dense bmm #5672 bmm(S, D)
• Convert dense matrix to sparse #2885 to_sparse()
• mul(sparse, dense) in pytorch #4456 Feature request: mul(Sparse, Dense) -> Sparse #2909 mul(Sparse, Dense)
• Feature request: sparse matrix max(axis) #4247 max(dim)
• torch.cuda.sparse.FloatTensor is not enabled. #8856 cuda()
• [Feature Request] Sparse-Dense elementwise Multiplication  #3158 broadcasting mul(S, D)

[li-roy]

I think norm is already implemented.

[weiyangfb]

@li-roy oh, you're right! Will remove it from todo

[weiyangfb]

also added TODO requests from @adefazio: ops and autograd support of

• mul(S, S) -> S
• mul(S, D) -> S
• dim-wise ops (softmax)

[realdoug]

@weiyangfb I took a pass at implementing narrow from your list. Posted a couple questions on the PR if you have a chance to look. Thanks!

[weiyangfb]

@realdoug Thanks for working on this feature! I will take a look!

[realdoug]

@weiyangfb is there a priority list at all for these ops? I think I can knock out a few more.

[weiyangfb]

@realdoug Thanks a lot for the great work! Let me take a look, stay tuned!

[Queuecumber]

I'd like to add einsum and mm(S,S)->S to the list if possible

[gongliyu]

Is it possible to add matmul(S, D) where S and D have rank > 2?

[rdv0011]

Is this topic related to conversion pytorch to ONNX?
I am facing the error when GraphEncoder::EncodeTensor() tries to call SparseTensorImpl::is_contiguous() which is not implemented. Please check the details here: onnx/onnx#2676

[rohanrajpal]

Hey

I'm trying to implement cosine similarity of two csr matrices in Pytorch

def awesome_cossim_top(A, B, ntop, lower_bound=0):
  # force A and B as a CSR matrix.
  # If they have already been CSR, there is no overhead
  A = A.tocsr()
  B = B.tocsr()
  M, _ = A.shape
  _, N = B.shape

  idx_dtype = np.int32

  nnz_max = M*ntop

  indptr = np.zeros(M+1, dtype=idx_dtype)
  indices = np.zeros(nnz_max, dtype=idx_dtype)
  data = np.zeros(nnz_max, dtype=A.dtype)

  ct.sparse_dot_topn(
      M, N, np.asarray(A.indptr, dtype=idx_dtype),
      np.asarray(A.indices, dtype=idx_dtype),
      A.data,
      np.asarray(B.indptr, dtype=idx_dtype),
      np.asarray(B.indices, dtype=idx_dtype),
      B.data,
      ntop,
      lower_bound,
      indptr, indices, data)

  return csr_matrix((data,indices,indptr),shape=(M,N))

So far I've reached

def csr_to_coo(X):
  X = X.tocoo()
  return torch.sparse.LongTensor(torch.LongTensor([X.row.tolist(), X.col.tolist()]),
                              torch.LongTensor(X.data.astype(np.int32)))

def cosine_distance(x1, x2=None, eps=1e-8):
    return 1 - torch.sparse.mm(x1, x2)

But if I run the code below

from scipy.sparse import csr_matrix

Acsr = csr_matrix([[1, 2, 0], [0, 0, 3], [4, 0, 5]])

print(cosine_distance(csr_to_coo(Acsr),csr_to_coo(Acsr)))

it throws the following error

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-16-eeed50549737> in <module>()
     21   return torch.sparse.LongTensor(torch.LongTensor([X.row.tolist(), X.col.tolist()]),
     22                               torch.LongTensor(X.data.astype(np.int32)))
---> 23 print(cosine_distance(csr_to_coo(Acsr),csr_to_coo(Acsr)))

1 frames
<ipython-input-16-eeed50549737> in cosine_distance(x1, x2, eps)
      5     # w2 = w1 if x2 is x1 else x2.norm(p=2, dim=1, keepdim=True)
      6     # return 1 - torch.sparse.mm(x1, x2.t()) / (w1 * w2.t()).clamp(min=eps)
----> 7     return 1 - torch.sparse.mm(x1, x2)
      8 
      9 Acsr = csr_matrix([[1, 2, 0], [0, 0, 3], [4, 0, 5]])

/usr/local/lib/python3.6/dist-packages/torch/sparse/__init__.py in mm(mat1, mat2)
     66                size=(2, 3), nnz=6, layout=torch.sparse_coo)
     67     """
---> 68     return torch._sparse_mm(mat1, mat2)
     69 
     70 

RuntimeError: sparse tensors do not have strides

From the ToDo I can see that (sparse, sparse) multiplication is supported, am I doing something wrong or this isn't supported yet?

[ezyang]

^ @aocsa @pearu maybe you'd be able to answer this?

[aocsa]

@rohanrajpal, sparse-sparse matmul is not supported in current master branch. Current torch.sparse.mm(x1, x2) function support sparse-dense matmul. The sparse-sparse matmul is an WIP in this PR.

[rohanrajpal]

@rohanrajpal, sparse-sparse matmul is not supported in current master branch. Current torch.sparse.mm(x1, x2) function support sparse-dense matmul. The sparse-sparse matmul is an WIP in this PR.

Alright, thanks!

[tczhangzhi]

@weiyangfb, I am working on graph convolutional neural networks, and need to perform expm1 operations on sparse tensors. In this s #46914 PR, I implemented an expm1() with autograd. It would be nice if someone can help review it.

[vadimkantorov]

Also, it would be good to support torch.min / torch.amin.

These could then be used in absence of #22378