Repository: pytorch
Issue: Semantics of sparse operations clarification - Sparsity of the gradient with respect to a sparse tensor input
Labels: module: sparse,module: docs,module: autograd,triaged,enhancement,module: masked operators

### 🚀 The feature, motivation and pitch

In many applications involving sparse matrices, one may be interested in imposing the sparsity pattern and using the gradient with repect to the "non-zero" elements of the sparse tensor only.

As a motivating toy example, let us consider tridiagonal matrices and let L be a (loss) function mapping such matrices to a scalar value, it may be of interest to find:
T_opt = argmin_{T ∈ Tridiagonal} L(T)

Using autodiff in such a context should ideally be straighforward with PyTorch. At the moment, however, the semantics of sparse matrix operations when it comes to the gradient with respect to the sparse matrix are not made explicit in the [documentation](https://pytorch.org/docs/stable/sparse.html) and may thus not match user expectations.

After discussion in several issues (see e.g. https://github.com/pytorch/pytorch/issues/87358#issuecomment-1285755188), I realised that the default (implicit) semantics of sparse tensors in PyTorch is that the sparse representation only affects the way the data is stored. Results of any operation with a sparse matrix, including backwards, should be equivalent to that of the operation applied to a dense view of the sparse matrix.

While this "transparent storage layout" semantics (as referred to it in https://github.com/pytorch/maskedtensor/issues/71#issuecomment-1286305562) makes sense in some setups, it may not be the semantics expected from a linear algebra perspective. In the latter case, the expectation of the user may typically be that the gradient of a sparse op with repect to the sparse input should be a sparse tensor (let's call this "sparsity-preserving gradient" semantics).

Despite "transparent storage layout" semantics being the default in PyTorch, "sparsity-preserving gradient" semantics is already used in some cases, in particular in [`torch.sparse.mm`](https://pytorch.org/docs/stable/generated/torch.sparse.mm.html), see e.g. discussion in #41128 and https://github.com/pytorch/pytorch/issues/2389#issuecomment-1006557121 but also usage in https://github.com/rusty1s/pytorch_sparse and https://github.com/flaport/torch_sparse_solve.

It was argued in https://github.com/pytorch/pytorch/issues/87358#issuecomment-1285797058 that [`torch.masked`](https://pytorch.org/docs/master/masked.html) would provide semantics more aligned with "sparsity-preserving gradient", however, as per https://github.com/pytorch/maskedtensor/issues/71#issuecomment-1286305562,  it seems unclear how MaskedTensors should deal with non-trivial sparse ops such as sparse @ dense. Furthermore, as illustrated in #41128, "sparsity-preserving gradient" semantics can be particularly useful to manage memory requirements when dealing with large sparse matrices. It's unclear to me how MaskedTensors can provide optimal memory and computational usage in this case in comparison to directly using sparse tenors.

Given the above, it would be great to:
- Clarify the default semantics for sparse tensors in the documentation
- Provide options to support "sparsity-preserving gradient" in key operations (as done for `torch.sparse.mm`)

### Alternatives

Alternatives include:
- Rely on a third-party sparse matrix library supporting "sparsity-preserving gradient" semantics such as https://github.com/rusty1s/pytorch_sparse
- Expand generic third-party library for abstract linear operators such as https://github.com/cornellius-gp/linear_operator to take advantage of efficient sparse operations
- Only work on small-scale problems (to ensure the dense gradient matrices can fit in memory) for which "sparsity-preserving gradient" can be achieved in two steps: 1) dense gradient; 2) subsample

### Additional context

Related issues: #87358 #41128 https://github.com/pytorch/maskedtensor/issues/71 #9674 #86963

cc @nikitaved @pearu @cpuhrsch @amjames @bhosmer @svekars @carljparker @ezyang @albanD @zou3519 @gqchen @soulitzer @Lezcano @Varal7