Create a MaskedTensor from a single sparse tensor

[tvercaut]

🚀 The feature, motivation and pitch

Sparse tensors already encode a masking patterns. It would be convenient to be able to consider a sparse matrix as a MaskedTensor where the non-zero indices of the sparse matrix would be considered as a mask of True value.

Alternatives

Duplicating the input sparse matrix is doable but confusing and requires additional storage for the mask.

Here is code I couldn't test (see #70 ):

t_csr # some sparse CSR matrix

# Create mask from CSR
mask_csr = torch.sparse_csr_tensor(t_csr.crow_indices(), t_csr.col_indices(), torch.ones_like(t_csr.values(),dtype=torch.bool), dtype=torch.bool)

# Create MaskedTensor
t_mt = maskedtensor.masked_tensor(t_csr, mask_csr)

Additional context

This span from https://github.com/pytorch/pytorch/issues/87358#issuecomment-1285797058

[cpuhrsch]

You could also use .bool() I think.

>>> t = torch.randn(3, 5).relu()
>>> t_csr = t.to_sparse_csr()
>>> mt = torch.masked.masked_tensor(t_csr, t_csr.bool())
>>> t
tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.5402, 0.0000, 1.9094, 0.4668, 0.0000],
        [0.5780, 0.5083, 0.0000, 0.0000, 0.0000]])
>>> mt
MaskedTensor(
  [
    [      --,       --,       --,       --,       --],
    [  0.5402,       --,   1.9094,   0.4668,       --],
    [  0.5780,   0.5083,       --,       --,       --]
  ]
)

Is this what you have in mind?

[cpuhrsch]

Also, maskedtensor is now in core under torch.masked. We just landed the docs update, but it might be slow to populate. We should probably update this repository to mention this as well and deprecate the pypi packages. cc @george-qi

[tvercaut]

Thanks. If the sparse matrix t_csr contains values that are equal to 0 in t_csr.values() I guess this would be different.

[tvercaut]

Out of scope for this issue but not sure where to ask otherwise.

Also, maskedtensor is now in core under torch.masked. We just landed the docs update, but it might be slow to populate. We should probably update this repository to mention this as well and deprecate the pypi packages.

Does this mean that othe bug reports / feature requests should be made on the core repo?

Also, really out of scope now, I am a bit unclear as to what the scope of maskedtensor should be. https://github.com/pytorch/pytorch/issues/87358#issuecomment-1285797058 seems to indicate it should become a drop-in replacement for sparse matrices by endowing it with clear semantics when it comes to gradients with respect to non-zero entries. However, the sparse ops listed here https://pytorch.org/docs/stable/sparse.html#supported-linear-algebra-operations don't seem to be supported by maskedtensor. For example is the plan to get sparse @ dense matrix multiplication supported in maskedtensor?

[cpuhrsch]

@tvercaut - To answer your questions in order

1. t_csr might have uncompressed zeros.

You're right that t_csr could have uncompressed zeros. We have temporarily allowed ones_like to return what you want, but it's breaking the "transparent storage layout" rule so we'll likely change it again soon (torch.sparse is still in beta (see definition)).

>>> t_csr
tensor(crow_indices=tensor([ 0,  2,  6, 10, 14]),
       col_indices=tensor([1, 3, 1, 2, 3, 4, 0, 1, 2, 4, 0, 2, 3, 4]),
       values=tensor([1.2325, 1.2773, 0.4609, 0.0000, 0.0000, 0.6036, 0.0000,
                      1.3118, 1.0352, 0.0000, 1.0767, 0.0000, 2.0717, 0.0000]),
       size=(4, 5), nnz=14, layout=torch.sparse_csr)
>>> torch.masked.masked_tensor(t_csr, torch.ones_like(t_csr).bool())
MaskedTensor(
  [
    [      --,   1.2325,       --,   1.2773,       --],
    [      --,   0.4609,   0.0000,   0.0000,   0.6036],
    [  0.0000,   1.3118,   1.0352,       --,   0.0000],
    [  1.0767,       --,   0.0000,   2.0717,   0.0000]
  ]
)

Instead you could also use sparse_mask together with a dense Tensor created via ones_likes.

2. MaskedTensor in core and GH issues

Now that it's in core it'd be better to open GH issues over there. This repo will likely soon go away / get deprecated.

3. Scope of MaskedTensor

It is an extension of torch Tensors, not unlike NumPy's masked arrays, that lets you have Tensors whose logical elements don't have a value (but do still have an index). It is a tuple (data, mask) and acts like a regular torch.Tensor if the mask is full. You can use sparse storage formats to compress repeat undefined elements. In particular by compressing the mask of a MaskedTensor you end up with indices that can be shared with the data. There are many applications for this kind of extension, one of them is pruning and sparsity aware training etc. etc.

sparse @ dense matmul is difficult to define for MaskedTensor, because we currently enforce that for binary operations the masks must match. We have not yet made a decision as to what should happen if you add an element without a value to an element with a value. We do however provide APIs to create your own behavior here (via .mask() and .data()). Further, for sparse @ dense matmul users often want their masked out elements to mean "zero". For that it's better to convert the MT to a regular Tensor with masked out elements filled in with, well, zeros.

[tvercaut]

Thanks. Very helpful. Re 3, I created a new feature request here: https://github.com/pytorch/pytorch/issues/87448