Conversion from strided to batched sparse compressed tensor with a non-constant number of zeros in batches fails
Labels: module: sparse, triaged


## Issue description

As in the title.

The issue is created to discuss various approaches to supporting the strided-to-sparse-compressed conversion for the cases where the number of zeros in different batches is not equal.

## Code example

Consider the following batched tensor of 2-by-2 tensors:
```
x = [[[0, 1],  # batch 1
      [2, 0]],
     [[3, 0],  # batch 2
      [0, 4]],
    ]
```
that can be represented as a batched CSR tensor:
```python
>>> torch.tensor(x).to_sparse_csr()
tensor(crow_indices=tensor([[0, 1, 2],
                            [0, 1, 2]]),
       col_indices=tensor([[1, 0],
                           [0, 1]]),
       values=tensor([[1, 2],
                      [3, 4]]), size=(2, 2, 2), nnz=2, layout=torch.sparse_csr)
```
because both batches have equal number zeros: 2.

Next, consider a batched tensor with an unequal number of zeros in batches:
```
y = [[[0, 1],  # batch 1
      [2, 9]],
     [[3, 0],  # batch 2
      [0, 4]],
    ]
```
that currently cannot be represented as a batched CSR tensor:
```python
>>> torch.tensor(y).to_sparse_csr()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Expect the same number of specified elements per batch.
```
because the number of zeros in batches is different: 1 and 2, respectively.

## Discussion

The following approaches exist to create a batched CSR tensor from batches having unequal numbers of zeros.

### Approach 1: allow materialization of certain zeros

Notice that in the conversion of a strided to a CSR tensor, non-zero elements and specified elements are considered as synonyms. If we relax this condition and allow certain zero elements to become specified elements for the CSR representation, the example tensor `y` defined above can be represented as a batched CSR tensor. In fact, there exist many such representations, for example:
```python
>>> torch.sparse_csr_tensor([[0, 1, 3], [0, 1, 3]], [[1, 0, 1], [0, 0, 1]], [[1, 2, 9], [3, 0, 4]]).to_dense()
tensor([[[0, 1],
         [2, 9]],
        [[3, 0],
         [0, 4]]])
>>> torch.sparse_csr_tensor([[0, 1, 3], [0, 2, 3]], [[1, 0, 1], [0, 1, 1]], [[1, 2, 9], [3, 0, 4]]).to_dense()
tensor([[[0, 1],
         [2, 9]],
        [[3, 0],
         [0, 4]]])
>>> torch.sparse_csr_tensor([[0, 2, 4], [0, 2, 4]], [[0, 1, 0, 1], [0, 1, 0, 1]], [[0, 1, 2, 9], [3, 0, 0, 4]]).to_dense()
tensor([[[0, 1],
         [2, 9]],
        [[3, 0],
         [0, 4]]])
```
that differ in the choice of materialized zeros.

Pros:
- solves the issue using the existing sparse compressed tensor implementation

Cons:
- requires introducing a complex and non-parallelizable strided->sparse compressed conversion algorithm that materializes zeros (non-optimal storage optimization) with the ambiguity of selecting materialized zeros (provisioning of a mask would resolve the ambiguity).
- batches with smaller NSE have the same memory usage as the batches with the largest NSE
- the maximum NSE in batches can be larger than the number of non-zeros in the batch of a minimum number of zeros

### Approach 2: allow a variable number of specified elements in batches

A prototype of this approach is implemented at https://github.com/pytorch/pytorch/pull/84843

The example tensor `y` defined above can be represented as a batched CSR tensor uniquely:
```python
>>> z = torch.tensor(y).to_sparse_csr()
>>> z
tensor(crow_indices=tensor([[0, 1, 3],
                            [0, 1, 2]]),
       col_indices=tensor([[1, 0, 1],
                           [0, 1, 0]]),
       values=tensor([[1, 2, 9],
                      [3, 4, 0]]), size=(2, 2, 2), nnz=3,
       layout=torch.sparse_csr)
```
where each batch has the expected NSE count:
```python
>>> z[0]
tensor(crow_indices=tensor([0, 1, 3]),
       col_indices=tensor([1, 0, 1]),
       values=tensor([1, 2, 9]), size=(2, 2), nnz=3, layout=torch.sparse_csr)
>>> z[1]
tensor(crow_indices=tensor([0, 1, 2]),
       col_indices=tensor([0, 1]),
       values=tensor([3, 4]), size=(2, 2), nnz=2, layout=torch.sparse_csr)
```

Pros:
- solves the issue without explicitly materializing zeros
- the batched sparse compressed tensor representation is unique
- the strided->sparse compressed conversion algorithm is simple and parallelizable
- the maximum NSE is equal to the number of non-zeros in the batch of a minimum number of zeros
- the performance of `to_sparse_csr()` on CUDA tensors increased by 15%

Cons:
- requires relaxing the sparse compressed invariant: `compressed_index[..., -1] == nnz` becomes `compressed_index[..., -1] <= nnz`
- batches with smaller NSE have the same memory usage as the batches with the largest NSE (the optimal storage would require ragged tensor support)
- the indices and values of unused elements in batches with smaller NSE must be initialized to avoid confusing third-party libraries if the batched tensor has variable NSE (if the libraries treat batches as independent, there should be no confusion)

## System Info

- PyTorch version: main


cc @alexsamardzic @nikitaved @cpuhrsch @amjames @bhosmer