Repository: pytorch
Issue: Design API for accessing sparse tensor indices
Labels: module: sparse,triaged

### 🚀 The feature, motivation and pitch

A sparse tensor can be viewed as a pair of indices and values that define the set of specified tensor elements. PyTorch supports sparse tensors with COO and CSR formats and starting from https://github.com/pytorch/pytorch/pull/75831, also CSC, BSR, and BSC formats.

All mentioned sparse formats use the `.values()` method to retrieve the values of specified sparse tensor elements. However, the API for accessing the indices of sparse tensor elements lacks uniformity and completeness:
- the indices of COO tensors are provided via the `.indices()` method
- the indices of CSR and BSR tensors are provided via the `.crow_indices()` and `.col_indices()` methods
- the indices of CSC and BSC tensors can be accessed via the `.crow_indices()` and `.col_indices()` methods as a temporary defect. It is a defect because `crow_indices()` returns the compressed column indices and `col_indices()` returns the row indices.

The motivation for creating this issue is
1. to resolve the temporary defect of accessing the indices of CSC and BSC tensors
2. to design a more uniform API for accessing sparse tensor indices of any sparse format

In the following, a number of possible solutions will be outlined for discussion.

## ~Proposal 1~ A fix to the CSC/BSC defect.

To address the reported defect of CSC and BSC tensors, a natural extension to the current convention of accessing indices is to introduce new methods `.ccol_indices()` and `.row_indices()` that will return the compressed row indices and column indices, respectively.

In addition, the methods `.crow_indices()` and `.col_indices()` will be supported only for sparse CSR and BSR tensors. The methods `.ccol_indices()` and `.row_indices()` will be supported only for sparse CSC and BSC tensors.

This option is actionable per https://github.com/pytorch/pytorch/issues/76638#issuecomment-1115405207

## :+1: Proposal 2

To reduce the number of indices related methods within the class of compressed tensors (CSR, CSC, BSR, BSC), another option is to introduce new methods `.compressed_indices()` and `.plain_indices()` that will return compressed row indices and column indices, respectively, for CSR and BSR tensors, and compressed column indices and row indices for CSC and BSC tensors.

In addition, the methods `.crow_indices()` and `.col_indices()` would be deprecated.

## :tada: Proposal 3

To reduce the number of the related methods even more and to unify the API in between all sparse formats, the third option is to reuse the existing `.indices()` method for retrieving the indices data for all compressed tensors, that is, for CSR and BSR tensors, it would return a 2-tuple that first item contains compressed row indices and the second item contains column indices. Similar, for CSC and BSR tensors, the `.indices()` would return a tuple of compressed column indices and row indices.

In addition, the methods `.crow_indices()` and `.col_indices()` would be deprecated.

In general, for any new sparse tensor format, the `.indices()` would return an object that contains indices data in a form that is most natural for the given format.

## :rocket: Proposal 4

From https://github.com/pytorch/pytorch/issues/76638#issuecomment-1115398434: Introduce optional `dim` option to `.indices()` method so that
- for COO tensors, we would have `coo.indices(dim=d) == coo.indices()[d]`
- for CSR/... tensors, `csr.indices(dim=d) == (csr.crow_indices() if d == 0 else csr.col_indices())`
- as a possible extension, for strided tensors, `dense.indices(dim=...)` returns indices tensor equivalent to `dense.to_sparse_coo().indices(dim=...)`.

## :smile: Proposal 5

From https://github.com/pytorch/pytorch/issues/76638#issuecomment-1115398434: Introduce `Indices` data structure such that
- for COO tensors, `.indices()` has the current behaviour, that is, return a tensor of indices
- for CSR/... tensors, `csr.indices()` returns `Indices` instance such that `torch.tensor(csr.indices()) == csr.to_sparse_coo().indices()` but internally, the row indices are compressed within the `Indices` object.

### Alternatives

An alternative is to copy the sparse tensor indices API from `scipy.sparse` where [sparse CSR tensor indices](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_matrix.html), for instance, can be accessed via attributes `.indptr` and `.indices` that correspond to compressed row indices and column indices, respectively. Clearly, the `scipy.sparse` API is incompatible with the corresponding API in PyTorch (e.g. compare `indices` being a column indices specific attribute in SciPy vs being an all-dimensions indices specific method in PyTorch).


### Additional context

_No response_

cc @nikitaved @pearu @cpuhrsch @IvanYashchuk 

## A poll for the design proposal

To vote for a proposal, react with the corresponding emoji to this comment. Multiple reactions are OK.