I’m not sure what the best course of action here is, but let me explain how the current sparse tensor representation works.
The current sparse representation (http://pytorch.org/docs/sparse.html 57) supports hybrid sparse tensors, where you can say that the first n dimensions are sparse, and the rest are dense; e.g., if you have a 3D tensor which only specifies a few 2D matrices in the stack. But this isn’t really what you are looking for: if I understand correctly, what you want is the first dimension to be dense, and then the second and third dimensions sparse (but with the same sparsity pattern.)
One thing you could do is transpose your input so that the sparse dimensions are at the beginning (so you have a 3D tensor which is a sparse 2D “matrix” whose entries are vectors, each entry corresponding to one of the entries from the batched 2D matrix.) This doesn’t seem completely awful to me, but I don’t have a well-developed aesthetic sense for this sort of thing yet.