Is there a canonical way to exploit sparsity in batch operations torch.bmm() and torch.baddmm() yet? Interested mainly in sparse -> dense (dense -> sparse is also interesting).

If I have a batch of sparse input matrices, and a dense batch of matrices :

bmm is currently not implemented for torch.sparse.* modules. If you can store it in a list, you can simply do torch.mm(mat1_i, mat2). If your matricies are extremely sparse, this should be pretty good

Also, torch.smm(mat1_i, mat2) is also implemented, for sparse * dense -> sparse operations.

Thanks for the tips If my sparse matrices are in a list, do you mean something like this? :

import torch
import torch.autograd as ann
mat2 = ann.Variable(torch.rand(4, 5, 6), requires_grad=True)
mats = [ann.Variable(torch.zeros(4, 3, 5), requires_grad=True) for _ in range(3)]
for i in range(len(mats)):
result = torch.bmm(mats[i], mat2)
print result.size()

import torch
x = torch.rand(5,6)
# Sparse matrix of (0, 1) = 1; (2, 1) = 2, (3, 4) = 3
sparse = torch.sparse.FloatTensor(
torch.LongTensor([[0, 2, 3], [1, 1, 4]]), # Indicies
torch.FloatTensor([1, 2, 3])) # values
print(x)
print(sparse.to_dense())
print(torch.mm(sparse, x))
# This won't actually save space or compute, since it's so dense,
# but it will be a sparse tensor representation.
print(torch.smm(sparse, x))

Simply construct a list of your sparse tensors, and loop over them to do the batch mm.

Seems strange to me as you don’t define the sizes of the sparse matrix - it seems to arbitrarily pick the indices of the corner value as the size. What is the logic here?