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)