Repository: pytorch
Issue: Resize/reshape of sparse compressed tensors - design
Labels: module: sparse,triaged

## Issue description

Currently, resizing sparse compressed tensors works correctly only in the CSR layout and non-hybrid case. Other cases are silently broken and the resize operation produces invalid tensors.

The issue is relevant for designing the reshape operation on sparse compressed tensors as well.

## Code example

```python
>>> import torch
>>> s = torch.tensor([[1, 2], [3, 4]]).to_sparse_csc()
>>> torch._validate_sparse_csc_tensor_args(s.ccol_indices(), s.row_indices(), s.values(), s.shape)
>>> s.resize_(2, 4)
tensor(ccol_indices=tensor([0, 2, 4]),
       row_indices=tensor([0, 1, 0, 1]),
       values=tensor([1, 3, 2, 4]), size=(2, 4), nnz=4,
       layout=torch.sparse_csc)
>>> torch._validate_sparse_csc_tensor_args(s.ccol_indices(), s.row_indices(), s.values(), s.shape)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: ccol_indices.shape[-1] must be equal to the number of columns + 1 (=5), but got 3
```

- PyTorch version: master

## Design

The size of a compressed sparse tensor has the following general form
```python
batchsize + (ncompressed_indices * blocksize[0], nplain_indices * blocksize[1]) + densesize
```
(see https://pearu.github.io/bsr_tensor_invariants.html#shape-and-strides-invariants-1) where for CSR/CSC tensors `blocksize[0] == blocksize[1] == 1` by definition.

For the attributes of a compressed sparse tensor we have the following relations:
```python
compressed_indices.shape == batchsize + (ncompressed_indices + 1,)
plain_indices.shape == batchsize + (nnz,)
values.shape == batchsize + (nnz,) + blocksize + densesize
0 <= compressed_indices[..., i] - compressed_indices[..., i-1] <= nplain_indices
plain_indices.max() < nplain_indices
```
where for CSR/CSC tensors `blocksize == ()` by definition and `...` corresponds to batch dimensions.

### Resizing batch and dense dimensions

Resizing batch and dense dimensions is straightforward:
```python
compressed_indices.resize_(new_batchsize + ...)
plain_indices.resize_(new_batchsize + ...)
values.resize_(new_batchsize + ...)
size = new_batchsize + ...
```
and
```python
values.resize_(... + new_densesize)
size = ... + new_densesize
```
respectively. `...` correspond to sparse and dense dimensions, and batch or sparse dimensions, respectively.

### Ambiguity in resizing sparse dimensions and a solution

Resizing compressed and plain dimensions requires re-computing both compressed and plain dimensions. In the case of blocked formats, the resize is ambiguous in general as the system of equations:
```python
ncompressed_indices * blocksize[0] == new_compressed_size
nplain_indices * blocksize[1] == new_plain_size
```
may not have a solution at all (that would lead to an exception) or have multiple solutions. In case of multiple solutions, conservation of `blocksize` is preferred, that is,
```python
new_ncompressed_indices = new_compressed_size / blocksize[0]
new_nplain_indices = new_plain_size / blocksize[1]
```
if new sizes are greater or equal than the corresponding blocksizes, otherwise, new blocksizes will be re-computed by maximizing these.

### Ambiguity in changing dimensionality and a solution

While the total dimensionality is allowed to change in the resize operation, in general, specifying a new dimensionality (when `len(new_shape) != len(tensor.shape)`) is ambiguous without specifying whether the newly specified dimensions correspond to batch, sparse, or dense dimensions.

To resolve these ambiguities, we distinguish non-hybrid and hybrid shapes.
A non-hybrid shape (`new_shape`) is defined by
```python
2 <= len(new_shape) <= len(tensor.shape) - tensor.dense_dim()
```
that is, the last two items in `new_shape` correspond to sparse dimensions and possibly non-zero dense dimensions will be dropped.
`len(new_shape) < 2` is disallowed by sparse compressed tensor definition and would lead to an exception.

A hybrid shape is defined by
```python
len(tensor.shape) - tensor.dense_dim() < len(new_shape)
```
that is, batched and sparse dimensionalities are preserved in the shape change (but the corresponding dimensions are allowed to change).

Limitations: one cannot introduce new batch dimensions with this approach, only changing the existing batch dimensions is possible.

### Validity of the resize result

When resizing involves changing sparse dimensions (compressed or plain), the corresponding indices attributes may need to be updated as well as the indices' compression result depends on the sparse dimensions.

### Questions

Does this make sense?

cc @nikitaved @pearu @cpuhrsch @amjames