Efficient and robust calculation of diag(sparse @ diag @ sparse)
Labels: module: sparse, feature, triaged

### 🚀 The feature, motivation and pitch

This is a feature request for improving sparse tensors. I'm not really expecting the issue to be resolved any time soon, but I've written about my use case and where I have been having problems in case it is helpful for those who work on pytorch. (Or maybe I'm using sparse tensors in the wrong way!)

As part of a model I'm implementing, I'd like to be able to compute $\text{diag}\left(\mathbf{A} \mathbf{B} \mathbf{A}^T\right)$ efficiently, where $\mathbf{B}$ is a diagonal matrix, and $\mathbf{A}$ is sparse (and both are square matrices). In theory I can compute this without using much memory, but I've been struggling to get it working properly in pytorch.

There is a nice way to do this if you are willing to store $\mathbf{A}$ as a dense tensor, namely `torch.sum((A * B_diag) * A,-1)`, but unfortunately this won't scale if $\mathbf{A}$ is a very big matrix. So ideally I'd use the `torch.sparse` API instead!

I found that it is possible to do something roughly similar to this using `torch.sparse` (illustrated below - [^1]), but it doesn't quite work properly - backprop doesn't work.

There is a workaround - [^2] - representing $\mathbf{B}$ as a sparse tensor and performing 2 sparse mat muls, but it is not very robust. It would be nice to have a `torch.sparse.diagonal` function for this!

[^1] - torch sparse implementation backprop doesn't work

```python
import torch as t
t.manual_seed(0)
B_diag = t.randn((10,), requires_grad=True)
A_dense = t.randn((10, 10), requires_grad=False)

naive_calc = t.diagonal(A_dense @ t.diag_embed(B_diag) @ A_dense.t())
dense_calc = t.sum(A_dense * B_diag * A_dense, -1)
assert (naive_calc - dense_calc).abs().max() < 1e-6

A_sparse = A_dense.to_sparse()
sparse_calc = t.sparse.sum(A_sparse * B_diag * A_sparse, -1).to_dense()
assert (naive_calc - sparse_calc).abs().max() < 1e-6 # this actually works! nice

# but i can't backprop through sparse_calc
# as an example
x1 = t.randn((10,), requires_grad = True)
(x1 - sparse_calc).sum().backward() # ERROR
```

### Alternatives

[^2] - In seeking an workaround, we can convert $\mathbf{B}$ to a sparse tensor, and compute 2 sparse matmuls. But it is not robust to having zeros on the diagonal

```python
ABAT = t.sparse.mm(t.sparse.mm(A_sparse, t.diag_embed(B_diag).to_sparse()), A_sparse.transpose(-2, -1))

# to get the diagonal of this sparse matrix, there is no t.sparse.diagonal function, so we can use a mask
eye_mask = t.eye(10, requires_grad=False).to_sparse() # could do this more efficiently if it were necessary
diag_ABAT = (ABAT * eye_mask).coalesce().values()
assert (diag_ABAT - naive_calc).abs().max() < 1e-6

# note this only works because every element on the non-diag is non-zero
ABAT_dense = ABAT.to_dense(); ABAT_dense[0,0] = 0.
ABAT = ABAT_dense.to_sparse()
diag_ABAT = (ABAT * eye_mask).coalesce().values()
assert (diag_ABAT - naive_calc).abs().max() < 1e-6 # ERROR!
```

### Additional context

torch version:  2.0.1
gpu: NVIDIA GeForce RTX 2080 Ti
cpu: Intel(R) Xeon(R) Silver 4112 CPU @ 2.60GHz

cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer