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:

>>> 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:

>>> 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:

>>> 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 #84843

The example tensor y defined above can be represented as a batched CSR tensor uniquely:

>>> 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:

>>> 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