.. automodule:: torch.sparse

.. currentmodule:: torch

.. _sparse-docs:

torch.sparse
============

.. warning::

  The PyTorch API of sparse tensors is in beta and may change in the near future.
  We highly welcome feature requests, bug reports and general suggestions as GitHub issues.

Why and when to use sparsity
++++++++++++++++++++++++++++

By default PyTorch stores :class:`torch.Tensor` stores elements contiguously
physical memory. This leads to efficient implementations of various array
processing algorithms that require fast access to elements.

Now, some users might decide to represent data such as graph adjacency
matrices, pruned weights or points clouds by Tensors whose *elements are
mostly zero valued*. We recognize these are important applications and aim
to provide performance optimizations for these use cases via sparse storage formats.

Various sparse storage formats such as COO, CSR/CSC, LIL, etc. have been
developed over the years. While they differ in exact layouts, they all
compress data through efficient representation of zero valued elements.
We call the uncompressed values *specified* in contrast to *unspecified*,
compressed elements.

By compressing repeat zeros sparse storage formats aim to save memory
and computational resources on various CPUs and GPUs. Especially for high
degrees of sparsity or highly structured sparsity this can have significant
performance implications. As such sparse storage formats can be seen as a
performance optimization.

Like many other performance optimization sparse storage formats are not
always advantageous. When trying sparse formats for your use case
you might find your execution time to increase rather than decrease.

Please feel encouraged to open a GitHub issue if you analytically
expected to see a stark increase in performance but measured a
degradation instead. This helps us prioritize the implementation
of efficient kernels and wider performance optimizations.

We make it easy to try different sparsity layouts, and convert between them,
without being opinionated on what's best for your particular application.

Functionality overview
++++++++++++++++++++++

We want it to be straightforward to construct a sparse Tensor from a
given dense Tensor by providing conversion routines for each layout.

In the next example we convert a 2D Tensor with default dense (strided)
layout to a 2D Tensor backed by the COO memory layout. Only values and
indices of non-zero elements are stored in this case.

    >>> a = torch.tensor([[0, 2.], [3, 0]])
    >>> a.to_sparse()
    tensor(indices=tensor([[0, 1],
                           [1, 0]]),
           values=tensor([2., 3.]),
           size=(2, 2), nnz=2, layout=torch.sparse_coo)

PyTorch currently supports :ref:`COO<sparse-coo-docs>`, :ref:`CSR<sparse-csr-docs>`,
:ref:`CSC<sparse-csc-docs>`, :ref:`BSR<sparse-bsr-docs>`, and :ref:`BSC<sparse-bsc-docs>`.
Please see the references for more details.

Note that we provide slight generalizations of these formats.

Batching: Devices such as GPUs require batching for optimal performance and
thus we support batch dimensions.

We currently offer a very simple version of batching where each component of a sparse format
itself is batched. This also requires the same number of specified elements per batch entry.
In this example we construct a 3D (batched) CSR Tensor from a 3D dense Tensor.

    >>> t = torch.tensor([[[1., 0], [2., 3.]], [[4., 0], [5., 6.]]])
    >>> t.dim()
    3
    >>> t.to_sparse_csr()
    tensor(crow_indices=tensor([[0, 1, 3],
                                [0, 1, 3]]),
           col_indices=tensor([[0, 0, 1],
                               [0, 0, 1]]),
           values=tensor([[1., 2., 3.],
                          [4., 5., 6.]]), size=(2, 2, 2), nnz=3,
           layout=torch.sparse_csr)


Dense dimensions: On the other hand, some data such as Graph embeddings might be
better viewed as sparse collections of vectors instead of scalars.

In this example we create a 3D Hybrid COO Tensor with 2 sparse and 1 dense dimension
from a 3D strided Tensor. If an entire row in the 3D strided Tensor is zero, it is
not stored. If however any of the values in the row are non-zero, they are stored
entirely. This reduces the number of indices since we need one index one per row instead
of one per element. But it also increases the amount of storage for the values. Since
only rows that are *entirely* zero can be emitted and the presence of any non-zero
valued elements cause the entire row to be stored.

    >>> t = torch.tensor([[[0., 0], [1., 2.]], [[0., 0], [3., 4.]]])
    >>> t.to_sparse(sparse_dim=2)
    tensor(indices=tensor([[0, 1],
                           [1, 1]]),
           values=tensor([[1., 2.],
                          [3., 4.]]),
           size=(2, 2, 2), nnz=2, layout=torch.sparse_coo)


Operator overview
+++++++++++++++++

Fundamentally, operations on Tensor with sparse storage formats behave the same as
operations on Tensor with strided (or other) storage formats. The particularities of
storage, that is the physical layout of the data, influences the performance of
an operation but should not influence the semantics.


We are actively increasing operator coverage for sparse tensors. Users should not
expect support same level of support as for dense Tensors yet.
See our :ref:`operator<sparse-ops-docs>` documentation for a list.

    >>> b = torch.tensor([[0, 0, 1, 2, 3, 0], [4, 5, 0, 6, 0, 0]])
    >>> b_s = b.to_sparse_csr()
    >>> b_s.cos()
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
    RuntimeError: unsupported tensor layout: SparseCsr
    >>> b_s.sin()
    tensor(crow_indices=tensor([0, 3, 6]),
           col_indices=tensor([2, 3, 4, 0, 1, 3]),
           values=tensor([ 0.8415,  0.9093,  0.1411, -0.7568, -0.9589, -0.2794]),
           size=(2, 6), nnz=6, layout=torch.sparse_csr)

As shown in the example above, we don't support non-zero preserving unary
operators such as cos. The output of a non-zero preserving unary operation
will not be able to take advantage of sparse storage formats to the same
extent as the input and potentially result in a catastrophic increase in memory.
We instead rely on the user to explicitly convert to a dense Tensor first and
then run the operation.

    >>> b_s.to_dense().cos()
    tensor([[ 1.0000, -0.4161],
            [-0.9900,  1.0000]])

We are aware that some users want to ignore compressed zeros for operations such
as `cos` instead of preserving the exact semantics of the operation. For this we
can point to torch.masked and its MaskedTensor, which is in turn also backed and
powered by sparse storage formats and kernels.

Also note that, for now, the user doesn't have a choice of the output layout. For example,
adding a sparse Tensor to a regular strided Tensor results in a strided Tensor. Some
users might prefer for this to stay a sparse layout, because they know the result will
still be sufficiently sparse.

    >>> a + b.to_sparse()
    tensor([[0., 3.],
            [3., 0.]])

We acknowledge that access to kernels that can efficiently produce different output
layouts can be very useful. A subsequent operation might significantly benefit from
receiving a particular layout. We are working on an API to control the result layout
and recognize it is an important feature to plan a more optimal path of execution for
any given model.


.. _sparse-coo-docs:

Sparse COO tensors
++++++++++++++++++

PyTorch implements the so-called Coordinate format, or COO
format, as one of the storage formats for implementing sparse
tensors.  In COO format, the specified elements are stored as tuples
of element indices and the corresponding values. In particular,

  - the indices of specified elements are collected in ``indices``
    tensor of size ``(ndim, nse)`` and with element type
    ``torch.int64``,

  - the corresponding values are collected in ``values`` tensor of
    size ``(nse,)`` and with an arbitrary integer or floating point
    number element type,

where ``ndim`` is the dimensionality of the tensor and ``nse`` is the
number of specified elements.

.. note::

   The memory consumption of a sparse COO tensor is at least ``(ndim *
   8 + <size of element type in bytes>) * nse`` bytes (plus a constant
   overhead from storing other tensor data).

   The memory consumption of a strided tensor is at least
   ``product(<tensor shape>) * <size of element type in bytes>``.

   For example, the memory consumption of a 10 000 x 10 000 tensor
   with 100 000 non-zero 32-bit floating point numbers is at least
   ``(2 * 8 + 4) * 100 000 = 2 000 000`` bytes when using COO tensor
   layout and ``10 000 * 10 000 * 4 = 400 000 000`` bytes when using
   the default strided tensor layout. Notice the 200 fold memory
   saving from using the COO storage format.

Construction
------------

A sparse COO tensor can be constructed by providing the two tensors of
indices and values, as well as the size of the sparse tensor (when it
cannot be inferred from the indices and values tensors) to a function
:func:`torch.sparse_coo_tensor`.

Suppose we want to define a sparse tensor with the entry 3 at location
(0, 2), entry 4 at location (1, 0), and entry 5 at location (1, 2).
Unspecified elements are assumed to have the same value, fill value,
which is zero by default. We would then write:

    >>> i = [[0, 1, 1],
             [2, 0, 2]]
    >>> v =  [3, 4, 5]
    >>> s = torch.sparse_coo_tensor(i, v, (2, 3))
    >>> s
    tensor(indices=tensor([[0, 1, 1],
                           [2, 0, 2]]),
           values=tensor([3, 4, 5]),
           size=(2, 3), nnz=3, layout=torch.sparse_coo)
    >>> s.to_dense()
    tensor([[0, 0, 3],
            [4, 0, 5]])

Note that the input ``i`` is NOT a list of index tuples.  If you want
to write your indices this way, you should transpose before passing them to
the sparse constructor:

    >>> i = [[0, 2], [1, 0], [1, 2]]
    >>> v =  [3,      4,      5    ]
    >>> s = torch.sparse_coo_tensor(list(zip(*i)), v, (2, 3))
    >>> # Or another equivalent formulation to get s
    >>> s = torch.sparse_coo_tensor(torch.tensor(i).t(), v, (2, 3))
    >>> torch.sparse_coo_tensor(i.t(), v, torch.Size([2,3])).to_dense()
    tensor([[0, 0, 3],
            [4, 0, 5]])

An empty sparse COO tensor can be constructed by specifying its size
only:

    >>> torch.sparse_coo_tensor(size=(2, 3))
    tensor(indices=tensor([], size=(2, 0)),
           values=tensor([], size=(0,)),
           size=(2, 3), nnz=0, layout=torch.sparse_coo)

.. _sparse-hybrid-coo-docs:

Sparse hybrid COO tensors
-------------------------

PyTorch implements an extension of sparse tensors with scalar values
to sparse tensors with (contiguous) tensor values. Such tensors are
called hybrid tensors.

PyTorch hybrid COO tensor extends the sparse COO tensor by allowing
the ``values`` tensor to be a multi-dimensional tensor so that we
have:

  - the indices of specified elements are collected in ``indices``
    tensor of size ``(sparse_dims, nse)`` and with element type
    ``torch.int64``,

  - the corresponding (tensor) values are collected in ``values``
    tensor of size ``(nse, dense_dims)`` and with an arbitrary integer
    or floating point number element type.

.. note::

   We use (M + K)-dimensional tensor to denote a N-dimensional sparse
   hybrid tensor, where M and K are the numbers of sparse and dense
   dimensions, respectively, such that M + K == N holds.

Suppose we want to create a (2 + 1)-dimensional tensor with the entry
[3, 4] at location (0, 2), entry [5, 6] at location (1, 0), and entry
[7, 8] at location (1, 2). We would write

    >>> i = [[0, 1, 1],
             [2, 0, 2]]
    >>> v =  [[3, 4], [5, 6], [7, 8]]
    >>> s = torch.sparse_coo_tensor(i, v, (2, 3, 2))
    >>> s
    tensor(indices=tensor([[0, 1, 1],
                           [2, 0, 2]]),
           values=tensor([[3, 4],
                          [5, 6],
                          [7, 8]]),
           size=(2, 3, 2), nnz=3, layout=torch.sparse_coo)

    >>> s.to_dense()
    tensor([[[0, 0],
             [0, 0],
             [3, 4]],
            [[5, 6],
             [0, 0],
             [7, 8]]])

In general, if ``s`` is a sparse COO tensor and ``M =
s.sparse_dim()``, ``K = s.dense_dim()``, then we have the following
invariants:

  - ``M + K == len(s.shape) == s.ndim`` - dimensionality of a tensor
    is the sum of the number of sparse and dense dimensions,
  - ``s.indices().shape == (M, nse)`` - sparse indices are stored
    explicitly,
  - ``s.values().shape == (nse,) + s.shape[M : M + K]`` - the values
    of a hybrid tensor are K-dimensional tensors,
  - ``s.values().layout == torch.strided`` - values are stored as
    strided tensors.

.. note::

   Dense dimensions always follow sparse dimensions, that is, mixing
   of dense and sparse dimensions is not supported.

.. note::

   To be sure that a constructed sparse tensor has consistent indices,
   values, and size, the invariant checks can be enabled per tensor
   creation via ``check_invariants=True`` keyword argument, or
   globally using :class:`torch.sparse.check_sparse_tensor_invariants`
   context manager instance. By default, the sparse tensor invariants
   checks are disabled.

.. _sparse-uncoalesced-coo-docs:

Uncoalesced sparse COO tensors
------------------------------

PyTorch sparse COO tensor format permits sparse *uncoalesced* tensors,
where there may be duplicate coordinates in the indices; in this case,
the interpretation is that the value at that index is the sum of all
duplicate value entries. For example, one can specify multiple values,
``3`` and ``4``, for the same index ``1``, that leads to an 1-D
uncoalesced tensor:

    >>> i = [[1, 1]]
    >>> v =  [3, 4]
    >>> s=torch.sparse_coo_tensor(i, v, (3,))
    >>> s
    tensor(indices=tensor([[1, 1]]),
           values=tensor(  [3, 4]),
           size=(3,), nnz=2, layout=torch.sparse_coo)

while the coalescing process will accumulate the multi-valued elements
into a single value using summation:

    >>> s.coalesce()
    tensor(indices=tensor([[1]]),
           values=tensor([7]),
           size=(3,), nnz=1, layout=torch.sparse_coo)

In general, the output of :meth:`torch.Tensor.coalesce` method is a
sparse tensor with the following properties:

- the indices of specified tensor elements are unique,
- the indices are sorted in lexicographical order,
- :meth:`torch.Tensor.is_coalesced()` returns ``True``.

.. note::

  For the most part, you shouldn't have to care whether or not a
  sparse tensor is coalesced or not, as most operations will work
  identically given a sparse coalesced or uncoalesced tensor.

  However, some operations can be implemented more efficiently on
  uncoalesced tensors, and some on coalesced tensors.

  For instance, addition of sparse COO tensors is implemented by
  simply concatenating the indices and values tensors:

    >>> a = torch.sparse_coo_tensor([[1, 1]], [5, 6], (2,))
    >>> b = torch.sparse_coo_tensor([[0, 0]], [7, 8], (2,))
    >>> a + b
    tensor(indices=tensor([[0, 0, 1, 1]]),
           values=tensor([7, 8, 5, 6]),
           size=(2,), nnz=4, layout=torch.sparse_coo)

  If you repeatedly perform an operation that can produce duplicate
  entries (e.g., :func:`torch.Tensor.add`), you should occasionally
  coalesce your sparse tensors to prevent them from growing too large.

  On the other hand, the lexicographical ordering of indices can be
  advantageous for implementing algorithms that involve many element
  selection operations, such as slicing or matrix products.

Working with sparse COO tensors
-------------------------------

Let's consider the following example:

    >>> i = [[0, 1, 1],
             [2, 0, 2]]
    >>> v =  [[3, 4], [5, 6], [7, 8]]
    >>> s = torch.sparse_coo_tensor(i, v, (2, 3, 2))

As mentioned above, a sparse COO tensor is a :class:`torch.Tensor`
instance and to distinguish it from the `Tensor` instances that use
some other layout, on can use :attr:`torch.Tensor.is_sparse` or
:attr:`torch.Tensor.layout` properties:

    >>> isinstance(s, torch.Tensor)
    True
    >>> s.is_sparse
    True
    >>> s.layout == torch.sparse_coo
    True

The number of sparse and dense dimensions can be acquired using
methods :meth:`torch.Tensor.sparse_dim` and
:meth:`torch.Tensor.dense_dim`, respectively. For instance:

    >>> s.sparse_dim(), s.dense_dim()
    (2, 1)


If ``s`` is a sparse COO tensor then its COO format data can be
acquired using methods :meth:`torch.Tensor.indices()` and
:meth:`torch.Tensor.values()`.

.. note::

  Currently, one can acquire the COO format data only when the tensor
  instance is coalesced:

    >>> s.indices()
    RuntimeError: Cannot get indices on an uncoalesced tensor, please call .coalesce() first

  For acquiring the COO format data of an uncoalesced tensor, use
  :func:`torch.Tensor._values()` and :func:`torch.Tensor._indices()`:

    >>> s._indices()
    tensor([[0, 1, 1],
            [2, 0, 2]])

  .. See https://github.com/pytorch/pytorch/pull/45695 for a new API.

  .. warning::
    Calling :meth:`torch.Tensor._values()` will return a *detached* tensor.
    To track gradients, :meth:`torch.Tensor.coalesce().values()` must be
    used instead.

Constructing a new sparse COO tensor results a tensor that is not
coalesced:

    >>> s.is_coalesced()
    False

but one can construct a coalesced copy of a sparse COO tensor using
the :meth:`torch.Tensor.coalesce` method:

    >>> s2 = s.coalesce()
    >>> s2.indices()
    tensor([[0, 1, 1],
           [2, 0, 2]])

When working with uncoalesced sparse COO tensors, one must take into
an account the additive nature of uncoalesced data: the values of the
same indices are the terms of a sum that evaluation gives the value of
the corresponding tensor element. For example, the scalar
multiplication on a sparse uncoalesced tensor could be implemented by
multiplying all the uncoalesced values with the scalar because ``c *
(a + b) == c * a + c * b`` holds. However, any nonlinear operation,
say, a square root, cannot be implemented by applying the operation to
uncoalesced data because ``sqrt(a + b) == sqrt(a) + sqrt(b)`` does not
hold in general.

Slicing (with positive step) of a sparse COO tensor is supported only
for dense dimensions. Indexing is supported for both sparse and dense
dimensions:

    >>> s[1]
    tensor(indices=tensor([[0, 2]]),
           values=tensor([[5, 6],
                          [7, 8]]),
           size=(3, 2), nnz=2, layout=torch.sparse_coo)
    >>> s[1, 0, 1]
    tensor(6)
    >>> s[1, 0, 1:]
    tensor([6])


In PyTorch, the fill value of a sparse tensor cannot be specified
explicitly and is assumed to be zero in general. However, there exists
operations that may interpret the fill value differently. For
instance, :func:`torch.sparse.softmax` computes the softmax with the
assumption that the fill value is negative infinity.

.. See https://github.com/Quansight-Labs/rfcs/tree/pearu/rfc-fill-value/RFC-0004-sparse-fill-value for a new API

.. _sparse-compressed-docs:

Sparse Compressed Tensors
+++++++++++++++++++++++++

Sparse Compressed Tensors represents a class of sparse tensors that
have a common feature of compressing the indices of a certain dimension
using an encoding that enables certain optimizations on linear algebra
kernels of sparse compressed tensors. This encoding is based on the
`Compressed Sparse Row (CSR)`__ format that PyTorch sparse compressed
tensors extend with the support of sparse tensor batches, allowing
multi-dimensional tensor values, and storing sparse tensor values in
dense blocks.

__ https://en.wikipedia.org/wiki/Sparse_matrix#Compressed_sparse_row_(CSR,_CRS_or_Yale_format)

.. note::

   We use (B + M + K)-dimensional tensor to denote a N-dimensional
   sparse compressed hybrid tensor, where B, M, and K are the numbers
   of batch, sparse, and dense dimensions, respectively, such that
   ``B + M + K == N`` holds. The number of sparse dimensions for
   sparse compressed tensors is always two, ``M == 2``.

.. note::

   We say that an indices tensor ``compressed_indices`` uses CSR
   encoding if the following invariants are satisfied:

   - ``compressed_indices`` is a contiguous strided 32 or 64 bit
     integer tensor
   - ``compressed_indices`` shape is ``(*batchsize,
     compressed_dim_size + 1)`` where ``compressed_dim_size`` is the
     number of compressed dimensions (e.g. rows or columns)
   - ``compressed_indices[..., 0] == 0`` where ``...`` denotes batch
     indices
   - ``compressed_indices[..., compressed_dim_size] == nse`` where
     ``nse`` is the number of specified elements
   - ``0 <= compressed_indices[..., i] - compressed_indices[..., i -
     1] <= plain_dim_size`` for ``i=1, ..., compressed_dim_size``,
     where ``plain_dim_size`` is the number of plain dimensions
     (orthogonal to compressed dimensions, e.g. columns or rows).

   To be sure that a constructed sparse tensor has consistent indices,
   values, and size, the invariant checks can be enabled per tensor
   creation via ``check_invariants=True`` keyword argument, or
   globally using :class:`torch.sparse.check_sparse_tensor_invariants`
   context manager instance. By default, the sparse tensor invariants
   checks are disabled.

.. note::

   The generalization of sparse compressed layouts to N-dimensional
   tensors can lead to some confusion regarding the count of specified
   elements. When a sparse compressed tensor contains batch dimensions
   the number of specified elements will correspond to the number of such
   elements per-batch. When a sparse compressed tensor has dense dimensions
   the element considered is now the K-dimensional array. Also for block
   sparse compressed layouts the 2-D block is considered as the element
   being specified.  Take as an example a 3-dimensional block sparse
   tensor, with one batch dimension of length ``b``, and a block
   shape of ``p, q``. If this tensor has ``n`` specified elements, then
   in fact we have ``n`` blocks specified per batch. This tensor would
   have ``values`` with shape ``(b, n, p, q)``. This interpretation of the
   number of specified elements comes from all sparse compressed layouts
   being derived from the compression of a 2-dimensional matrix. Batch
   dimensions are treated as stacking of sparse matrices, dense dimensions
   change the meaning of the element from a simple scalar value to an
   array with its own dimensions.

.. _sparse-csr-docs:

Sparse CSR Tensor
-----------------

The primary advantage of the CSR format over the COO format is better
use of storage and much faster computation operations such as sparse
matrix-vector multiplication using MKL and MAGMA backends.

In the simplest case, a (0 + 2 + 0)-dimensional sparse CSR tensor
consists of three 1-D tensors: ``crow_indices``, ``col_indices`` and
``values``:

  - The ``crow_indices`` tensor consists of compressed row
    indices. This is a 1-D tensor of size ``nrows + 1`` (the number of
    rows plus 1). The last element of ``crow_indices`` is the number
    of specified elements, ``nse``. This tensor encodes the index in
    ``values`` and ``col_indices`` depending on where the given row
    starts. Each successive number in the tensor subtracted by the
    number before it denotes the number of elements in a given row.

  - The ``col_indices`` tensor contains the column indices of each
    element. This is a 1-D tensor of size ``nse``.

  - The ``values`` tensor contains the values of the CSR tensor
    elements. This is a 1-D tensor of size ``nse``.

.. note::

   The index tensors ``crow_indices`` and ``col_indices`` should have
   element type either ``torch.int64`` (default) or
   ``torch.int32``. If you want to use MKL-enabled matrix operations,
   use ``torch.int32``. This is as a result of the default linking of
   pytorch being with MKL LP64, which uses 32 bit integer indexing.

In the general case, the (B + 2 + K)-dimensional sparse CSR tensor
consists of two (B + 1)-dimensional index tensors ``crow_indices`` and
``col_indices``, and of (1 + K)-dimensional ``values`` tensor such
that

  - ``crow_indices.shape == (*batchsize, nrows + 1)``

  - ``col_indices.shape == (*batchsize, nse)``

  - ``values.shape == (nse, *densesize)``

while the shape of the sparse CSR tensor is ``(*batchsize, nrows,
ncols, *densesize)`` where ``len(batchsize) == B`` and
``len(densesize) == K``.

.. note::

   The batches of sparse CSR tensors are dependent: the number of
   specified elements in all batches must be the same. This somewhat
   artificial constraint allows efficient storage of the indices of
   different CSR batches.

.. note::

   The number of sparse and dense dimensions can be acquired using
   :meth:`torch.Tensor.sparse_dim` and :meth:`torch.Tensor.dense_dim`
   methods. The batch dimensions can be computed from the tensor
   shape: ``batchsize = tensor.shape[:-tensor.sparse_dim() -
   tensor.dense_dim()]``.

.. note::

   The memory consumption of a sparse CSR tensor is at least
   ``(nrows * 8 + (8 + <size of element type in bytes> *
   prod(densesize)) * nse) * prod(batchsize)`` bytes (plus a constant
   overhead from storing other tensor data).

   With the same example data of :ref:`the note in sparse COO format
   introduction<sparse-coo-docs>`, the memory consumption of a 10 000
   x 10 000 tensor with 100 000 non-zero 32-bit floating point numbers
   is at least ``(10000 * 8 + (8 + 4 * 1) * 100 000) * 1 = 1 280 000``
   bytes when using CSR tensor layout. Notice the 1.6 and 310 fold
   savings from using CSR storage format compared to using the COO and
   strided formats, respectively.

Construction of CSR tensors
'''''''''''''''''''''''''''

Sparse CSR tensors can be directly constructed by using the
:func:`torch.sparse_csr_tensor` function. The user must supply the row
and column indices and values tensors separately where the row indices
must be specified using the CSR compression encoding.  The ``size``
argument is optional and will be deduced from the ``crow_indices`` and
``col_indices`` if it is not present.

    >>> crow_indices = torch.tensor([0, 2, 4])
    >>> col_indices = torch.tensor([0, 1, 0, 1])
    >>> values = torch.tensor([1, 2, 3, 4])
    >>> csr = torch.sparse_csr_tensor(crow_indices, col_indices, values, dtype=torch.float64)
    >>> csr
    tensor(crow_indices=tensor([0, 2, 4]),
           col_indices=tensor([0, 1, 0, 1]),
           values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4,
           dtype=torch.float64)
    >>> csr.to_dense()
    tensor([[1., 2.],
            [3., 4.]], dtype=torch.float64)

.. note::

   The values of sparse dimensions in deduced ``size`` is computed
   from the size of ``crow_indices`` and the maximal index value in
   ``col_indices``. If the number of columns needs to be larger than
   in the deduced ``size`` then the ``size`` argument must be
   specified explicitly.

The simplest way of constructing a 2-D sparse CSR tensor from a
strided or sparse COO tensor is to use
:meth:`torch.Tensor.to_sparse_csr` method. Any zeros in the (strided)
tensor will be interpreted as missing values in the sparse tensor:

    >>> a = torch.tensor([[0, 0, 1, 0], [1, 2, 0, 0], [0, 0, 0, 0]], dtype=torch.float64)
    >>> sp = a.to_sparse_csr()
    >>> sp
    tensor(crow_indices=tensor([0, 1, 3, 3]),
          col_indices=tensor([2, 0, 1]),
          values=tensor([1., 1., 2.]), size=(3, 4), nnz=3, dtype=torch.float64)

CSR Tensor Operations
'''''''''''''''''''''

The sparse matrix-vector multiplication can be performed with the
:meth:`tensor.matmul` method. This is currently the only math operation
supported on CSR tensors.

    >>> vec = torch.randn(4, 1, dtype=torch.float64)
    >>> sp.matmul(vec)
    tensor([[0.9078],
            [1.3180],
            [0.0000]], dtype=torch.float64)

.. _sparse-csc-docs:

Sparse CSC Tensor
-----------------

The sparse CSC (Compressed Sparse Column) tensor format implements the
CSC format for storage of 2 dimensional tensors with an extension to
supporting batches of sparse CSC tensors and values being
multi-dimensional tensors.

.. note::

   Sparse CSC tensor is essentially a transpose of the sparse CSR
   tensor when the transposition is about swapping the sparse
   dimensions.

Similarly to :ref:`sparse CSR tensors <sparse-csr-docs>`, a sparse CSC
tensor consists of three tensors: ``ccol_indices``, ``row_indices``
and ``values``:

  - The ``ccol_indices`` tensor consists of compressed column
    indices. This is a (B + 1)-D tensor of shape ``(*batchsize, ncols + 1)``.
    The last element is the number of specified
    elements, ``nse``. This tensor encodes the index in ``values`` and
    ``row_indices`` depending on where the given column starts. Each
    successive number in the tensor subtracted by the number before it
    denotes the number of elements in a given column.

  - The ``row_indices`` tensor contains the row indices of each
    element. This is a (B + 1)-D tensor of shape ``(*batchsize, nse)``.

  - The ``values`` tensor contains the values of the CSC tensor
    elements. This is a (1 + K)-D tensor of shape ``(nse, *densesize)``.

Construction of CSC tensors
'''''''''''''''''''''''''''

Sparse CSC tensors can be directly constructed by using the
:func:`torch.sparse_csc_tensor` function. The user must supply the row
and column indices and values tensors separately where the column indices
must be specified using the CSR compression encoding.  The ``size``
argument is optional and will be deduced from the ``row_indices`` and
``ccol_indices`` tensors if it is not present.

    >>> ccol_indices = torch.tensor([0, 2, 4])
    >>> row_indices = torch.tensor([0, 1, 0, 1])
    >>> values = torch.tensor([1, 2, 3, 4])
    >>> csc = torch.sparse_csc_tensor(ccol_indices, row_indices, values, dtype=torch.float64)
    >>> csc
    tensor(ccol_indices=tensor([0, 2, 4]),
           row_indices=tensor([0, 1, 0, 1]),
           values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4,
           dtype=torch.float64, layout=torch.sparse_csc)
    >>> csc.to_dense()
    tensor([[1., 3.],
            [2., 4.]], dtype=torch.float64)

.. note::

   The sparse CSC tensor constructor function has the compressed
   column indices argument before the row indices argument.

The (0 + 2 + 0)-dimensional sparse CSC tensors can be constructed from
any two-dimensional tensor using :meth:`torch.Tensor.to_sparse_csc`
method. Any zeros in the (strided) tensor will be interpreted as
missing values in the sparse tensor:

    >>> a = torch.tensor([[0, 0, 1, 0], [1, 2, 0, 0], [0, 0, 0, 0]], dtype=torch.float64)
    >>> sp = a.to_sparse_csc()
    >>> sp
    tensor(ccol_indices=tensor([0, 1, 2, 3, 3]),
           row_indices=tensor([1, 1, 0]),
           values=tensor([1., 2., 1.]), size=(3, 4), nnz=3, dtype=torch.float64,
           layout=torch.sparse_csc)

.. _sparse-bsr-docs:

Sparse BSR Tensor
-----------------

The sparse BSR (Block compressed Sparse Row) tensor format implements the
BSR format for storage of two-dimensional tensors with an extension to
supporting batches of sparse BSR tensors and values being blocks of
multi-dimensional tensors.

A sparse BSR tensor consists of three tensors: ``crow_indices``,
``col_indices`` and ``values``:

  - The ``crow_indices`` tensor consists of compressed row
    indices. This is a (B + 1)-D tensor of shape ``(*batchsize,
    nrowblocks + 1)``.  The last element is the number of specified blocks,
    ``nse``. This tensor encodes the index in ``values`` and
    ``col_indices`` depending on where the given column block
    starts. Each successive number in the tensor subtracted by the
    number before it denotes the number of blocks in a given row.

  - The ``col_indices`` tensor contains the column block indices of each
    element. This is a (B + 1)-D tensor of shape ``(*batchsize,
    nse)``.

  - The ``values`` tensor contains the values of the sparse BSR tensor
    elements collected into two-dimensional blocks. This is a (1 + 2 +
    K)-D tensor of shape ``(nse, nrowblocks, ncolblocks,
    *densesize)``.

Construction of BSR tensors
'''''''''''''''''''''''''''

Sparse BSR tensors can be directly constructed by using the
:func:`torch.sparse_bsr_tensor` function. The user must supply the row
and column block indices and values tensors separately where the row block indices
must be specified using the CSR compression encoding.
The ``size`` argument is optional and will be deduced from the ``crow_indices`` and
``col_indices`` tensors if it is not present.

    >>> crow_indices = torch.tensor([0, 2, 4])
    >>> col_indices = torch.tensor([0, 1, 0, 1])
    >>> values = torch.tensor([[[0, 1, 2], [6, 7, 8]],
    ...                        [[3, 4, 5], [9, 10, 11]],
    ...                        [[12, 13, 14], [18, 19, 20]],
    ...                        [[15, 16, 17], [21, 22, 23]]])
    >>> bsr = torch.sparse_bsr_tensor(crow_indices, col_indices, values, dtype=torch.float64)
    >>> bsr
    tensor(crow_indices=tensor([0, 2, 4]),
           col_indices=tensor([0, 1, 0, 1]),
           values=tensor([[[ 0.,  1.,  2.],
                           [ 6.,  7.,  8.]],
                          [[ 3.,  4.,  5.],
                           [ 9., 10., 11.]],
                          [[12., 13., 14.],
                           [18., 19., 20.]],
                          [[15., 16., 17.],
                           [21., 22., 23.]]]),
           size=(4, 6), nnz=4, dtype=torch.float64, layout=torch.sparse_bsr)
    >>> bsr.to_dense()
    tensor([[ 0.,  1.,  2.,  3.,  4.,  5.],
            [ 6.,  7.,  8.,  9., 10., 11.],
            [12., 13., 14., 15., 16., 17.],
            [18., 19., 20., 21., 22., 23.]], dtype=torch.float64)

The (0 + 2 + 0)-dimensional sparse BSR tensors can be constructed from
any two-dimensional tensor using :meth:`torch.Tensor.to_sparse_bsr`
method that also requires the specification of the values block size:

    >>> dense = torch.tensor([[0, 1, 2, 3, 4, 5],
    ...                       [6, 7, 8, 9, 10, 11],
    ...                       [12, 13, 14, 15, 16, 17],
    ...                       [18, 19, 20, 21, 22, 23]])
    >>> bsr = dense.to_sparse_bsr(blocksize=(2, 3))
    >>> bsr
    tensor(crow_indices=tensor([0, 2, 4]),
           col_indices=tensor([0, 1, 0, 1]),
           values=tensor([[[ 0,  1,  2],
                           [ 6,  7,  8]],
                          [[ 3,  4,  5],
                           [ 9, 10, 11]],
                          [[12, 13, 14],
                           [18, 19, 20]],
                          [[15, 16, 17],
                           [21, 22, 23]]]), size=(4, 6), nnz=4,
           layout=torch.sparse_bsr)

.. _sparse-bsc-docs:

Sparse BSC Tensor
-----------------

The sparse BSC (Block compressed Sparse Column) tensor format implements the
BSC format for storage of two-dimensional tensors with an extension to
supporting batches of sparse BSC tensors and values being blocks of
multi-dimensional tensors.

A sparse BSC tensor consists of three tensors: ``ccol_indices``,
``row_indices`` and ``values``:

  - The ``ccol_indices`` tensor consists of compressed column
    indices. This is a (B + 1)-D tensor of shape ``(*batchsize,
    ncolblocks + 1)``.  The last element is the number of specified blocks,
    ``nse``. This tensor encodes the index in ``values`` and
    ``row_indices`` depending on where the given row block
    starts. Each successive number in the tensor subtracted by the
    number before it denotes the number of blocks in a given column.

  - The ``row_indices`` tensor contains the row block indices of each
    element. This is a (B + 1)-D tensor of shape ``(*batchsize,
    nse)``.

  - The ``values`` tensor contains the values of the sparse BSC tensor
    elements collected into two-dimensional blocks. This is a (1 + 2 +
    K)-D tensor of shape ``(nse, nrowblocks, ncolblocks,
    *densesize)``.

Construction of BSC tensors
'''''''''''''''''''''''''''

Sparse BSC tensors can be directly constructed by using the
:func:`torch.sparse_bsc_tensor` function. The user must supply the row
and column block indices and values tensors separately where the column block indices
must be specified using the CSR compression encoding.
The ``size`` argument is optional and will be deduced from the ``ccol_indices`` and
``row_indices`` tensors if it is not present.

    >>> ccol_indices = torch.tensor([0, 2, 4])
    >>> row_indices = torch.tensor([0, 1, 0, 1])
    >>> values = torch.tensor([[[0, 1, 2], [6, 7, 8]],
    ...                        [[3, 4, 5], [9, 10, 11]],
    ...                        [[12, 13, 14], [18, 19, 20]],
    ...                        [[15, 16, 17], [21, 22, 23]]])
    >>> bsc = torch.sparse_bsc_tensor(ccol_indices, row_indices, values, dtype=torch.float64)
    >>> bsc
    tensor(ccol_indices=tensor([0, 2, 4]),
           row_indices=tensor([0, 1, 0, 1]),
           values=tensor([[[ 0.,  1.,  2.],
                           [ 6.,  7.,  8.]],
                          [[ 3.,  4.,  5.],
                           [ 9., 10., 11.]],
                          [[12., 13., 14.],
                           [18., 19., 20.]],
                          [[15., 16., 17.],
                           [21., 22., 23.]]]), size=(4, 6), nnz=4,
           dtype=torch.float64, layout=torch.sparse_bsc)

Tools for working with sparse compressed tensors
------------------------------------------------

All sparse compressed tensors --- CSR, CSC, BSR, and BSC tensors ---
are conceptionally very similar in that their indices data is split
into two parts: so-called compressed indices that use the CSR
encoding, and so-called plain indices that are orthogonal to the
compressed indices. This allows various tools on these tensors to
share the same implementations that are parameterized by tensor
layout.

Construction of sparse compressed tensors
'''''''''''''''''''''''''''''''''''''''''

Sparse CSR, CSC, BSR, and CSC tensors can be constructed by using
:func:`torch.sparse_compressed_tensor` function that have the same
interface as the above discussed constructor functions
:func:`torch.sparse_csr_tensor`, :func:`torch.sparse_csc_tensor`,
:func:`torch.sparse_bsr_tensor`, and :func:`torch.sparse_bsc_tensor`,
respectively, but with an extra required ``layout`` argument. The
following example illustrates a method of constructing CSR and CSC
tensors using the same input data by specifying the corresponding
layout parameter to the :func:`torch.sparse_compressed_tensor`
function:

    >>> compressed_indices = torch.tensor([0, 2, 4])
    >>> plain_indices = torch.tensor([0, 1, 0, 1])
    >>> values = torch.tensor([1, 2, 3, 4])
    >>> csr = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, layout=torch.sparse_csr)
    >>> csr
    tensor(crow_indices=tensor([0, 2, 4]),
           col_indices=tensor([0, 1, 0, 1]),
           values=tensor([1, 2, 3, 4]), size=(2, 2), nnz=4,
           layout=torch.sparse_csr)
    >>> csc = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, layout=torch.sparse_csc)
    >>> csc
    tensor(ccol_indices=tensor([0, 2, 4]),
           row_indices=tensor([0, 1, 0, 1]),
           values=tensor([1, 2, 3, 4]), size=(2, 2), nnz=4,
           layout=torch.sparse_csc)
    >>> (csr.transpose(0, 1).to_dense() == csc.to_dense()).all()
    tensor(True)

.. _sparse-ops-docs:

Supported operations
+++++++++++++++++++++++++++++++++++

Linear Algebra operations
-------------------------

The following table summarizes supported Linear Algebra operations on
sparse matrices where the operands layouts may vary. Here
``T[layout]`` denotes a tensor with a given layout. Similarly,
``M[layout]`` denotes a matrix (2-D PyTorch tensor), and ``V[layout]``
denotes a vector (1-D PyTorch tensor). In addition, ``f`` denotes a
scalar (float or 0-D PyTorch tensor), ``*`` is element-wise
multiplication, and ``@`` is matrix multiplication.

.. csv-table::
   :header: "PyTorch operation", "Sparse grad?", "Layout signature"
   :widths: 20, 5, 60
   :delim: ;

   :func:`torch.mv`;no; ``M[sparse_coo] @ V[strided] -> V[strided]``
   :func:`torch.mv`;no; ``M[sparse_csr] @ V[strided] -> V[strided]``
   :func:`torch.matmul`; no; ``M[sparse_coo] @ M[strided] -> M[strided]``
   :func:`torch.matmul`; no; ``M[sparse_csr] @ M[strided] -> M[strided]``
   :func:`torch.mm`; no; ``M[sparse_coo] @ M[strided] -> M[strided]``
   :func:`torch.sparse.mm`; yes; ``M[sparse_coo] @ M[strided] -> M[strided]``
   :func:`torch.smm`; no; ``M[sparse_coo] @ M[strided] -> M[sparse_coo]``
   :func:`torch.hspmm`; no; ``M[sparse_coo] @ M[strided] -> M[hybrid sparse_coo]``
   :func:`torch.bmm`; no; ``T[sparse_coo] @ T[strided] -> T[strided]``
   :func:`torch.addmm`; no; ``f * M[strided] + f * (M[sparse_coo] @ M[strided]) -> M[strided]``
   :func:`torch.sparse.addmm`; yes; ``f * M[strided] + f * (M[sparse_coo] @ M[strided]) -> M[strided]``
   :func:`torch.sspaddmm`; no; ``f * M[sparse_coo] + f * (M[sparse_coo] @ M[strided]) -> M[sparse_coo]``
   :func:`torch.lobpcg`; no; ``GENEIG(M[sparse_coo]) -> M[strided], M[strided]``
   :func:`torch.pca_lowrank`; yes; ``PCA(M[sparse_coo]) -> M[strided], M[strided], M[strided]``
   :func:`torch.svd_lowrank`; yes; ``SVD(M[sparse_coo]) -> M[strided], M[strided], M[strided]``

where "Sparse grad?" column indicates if the PyTorch operation supports
backward with respect to sparse matrix argument. All PyTorch operations,
except :func:`torch.smm`, support backward with respect to strided
matrix arguments.

.. note::

   Currently, PyTorch does not support matrix multiplication with the
   layout signature ``M[strided] @ M[sparse_coo]``. However,
   applications can still compute this using the matrix relation ``D @
   S == (S.t() @ D.t()).t()``.

Tensor methods and sparse
-------------------------

The following Tensor methods are related to sparse tensors:

.. autosummary::
    :toctree: generated
    :nosignatures:

    Tensor.is_sparse
    Tensor.is_sparse_csr
    Tensor.dense_dim
    Tensor.sparse_dim
    Tensor.sparse_mask
    Tensor.to_sparse
    Tensor.to_sparse_coo
    Tensor.to_sparse_csr
    Tensor.to_sparse_csc
    Tensor.to_sparse_bsr
    Tensor.to_sparse_bsc
    Tensor.to_dense
    Tensor.values

The following Tensor methods are specific to sparse COO tensors:

.. autosummary::
    :toctree: generated
    :nosignatures:

    Tensor.coalesce
    Tensor.sparse_resize_
    Tensor.sparse_resize_and_clear_
    Tensor.is_coalesced
    Tensor.indices

The following methods are specific to :ref:`sparse CSR tensors <sparse-csr-docs>` and :ref:`sparse BSR tensors <sparse-bsr-docs>`:

.. autosummary::
    :toctree: generated
    :nosignatures:

    Tensor.crow_indices
    Tensor.col_indices

The following methods are specific to :ref:`sparse CSC tensors <sparse-csc-docs>` and :ref:`sparse BSC tensors <sparse-bsc-docs>`:

.. autosummary::
    :toctree: generated
    :nosignatures:

    Tensor.row_indices
    Tensor.ccol_indices

The following Tensor methods support sparse COO tensors:

:meth:`~torch.Tensor.add`
:meth:`~torch.Tensor.add_`
:meth:`~torch.Tensor.addmm`
:meth:`~torch.Tensor.addmm_`
:meth:`~torch.Tensor.any`
:meth:`~torch.Tensor.asin`
:meth:`~torch.Tensor.asin_`
:meth:`~torch.Tensor.arcsin`
:meth:`~torch.Tensor.arcsin_`
:meth:`~torch.Tensor.bmm`
:meth:`~torch.Tensor.clone`
:meth:`~torch.Tensor.deg2rad`
:meth:`~torch.Tensor.deg2rad_`
:meth:`~torch.Tensor.detach`
:meth:`~torch.Tensor.detach_`
:meth:`~torch.Tensor.dim`
:meth:`~torch.Tensor.div`
:meth:`~torch.Tensor.div_`
:meth:`~torch.Tensor.floor_divide`
:meth:`~torch.Tensor.floor_divide_`
:meth:`~torch.Tensor.get_device`
:meth:`~torch.Tensor.index_select`
:meth:`~torch.Tensor.isnan`
:meth:`~torch.Tensor.log1p`
:meth:`~torch.Tensor.log1p_`
:meth:`~torch.Tensor.mm`
:meth:`~torch.Tensor.mul`
:meth:`~torch.Tensor.mul_`
:meth:`~torch.Tensor.mv`
:meth:`~torch.Tensor.narrow_copy`
:meth:`~torch.Tensor.neg`
:meth:`~torch.Tensor.neg_`
:meth:`~torch.Tensor.negative`
:meth:`~torch.Tensor.negative_`
:meth:`~torch.Tensor.numel`
:meth:`~torch.Tensor.rad2deg`
:meth:`~torch.Tensor.rad2deg_`
:meth:`~torch.Tensor.resize_as_`
:meth:`~torch.Tensor.size`
:meth:`~torch.Tensor.pow`
:meth:`~torch.Tensor.sqrt`
:meth:`~torch.Tensor.square`
:meth:`~torch.Tensor.smm`
:meth:`~torch.Tensor.sspaddmm`
:meth:`~torch.Tensor.sub`
:meth:`~torch.Tensor.sub_`
:meth:`~torch.Tensor.t`
:meth:`~torch.Tensor.t_`
:meth:`~torch.Tensor.transpose`
:meth:`~torch.Tensor.transpose_`
:meth:`~torch.Tensor.zero_`

Torch functions specific to sparse Tensors
------------------------------------------

.. autosummary::
    :toctree: generated
    :nosignatures:

    sparse_coo_tensor
    sparse_csr_tensor
    sparse_csc_tensor
    sparse_bsr_tensor
    sparse_bsc_tensor
    sparse_compressed_tensor
    sparse.sum
    sparse.addmm
    sparse.sampled_addmm
    sparse.mm
    sspaddmm
    hspmm
    smm
    sparse.softmax
    sparse.log_softmax
    sparse.spdiags

Other functions
---------------

The following :mod:`torch` functions support sparse tensors:

:func:`~torch.cat`
:func:`~torch.dstack`
:func:`~torch.empty`
:func:`~torch.empty_like`
:func:`~torch.hstack`
:func:`~torch.index_select`
:func:`~torch.is_complex`
:func:`~torch.is_floating_point`
:func:`~torch.is_nonzero`
:func:`~torch.is_same_size`
:func:`~torch.is_signed`
:func:`~torch.is_tensor`
:func:`~torch.lobpcg`
:func:`~torch.mm`
:func:`~torch.native_norm`
:func:`~torch.pca_lowrank`
:func:`~torch.select`
:func:`~torch.stack`
:func:`~torch.svd_lowrank`
:func:`~torch.unsqueeze`
:func:`~torch.vstack`
:func:`~torch.zeros`
:func:`~torch.zeros_like`

To manage checking sparse tensor invariants, see:

.. autosummary::
    :toctree: generated
    :nosignatures:

    sparse.check_sparse_tensor_invariants

Unary functions
---------------

We aim to support all zero-preserving unary functions.

If you find that we are missing a zero-preserving unary function
that you need, please feel encouraged to open an issue for a feature request.
As always please kindly try the search function first before opening an issue.

The following operators currently support sparse COO/CSR/CSC/BSR/CSR tensor inputs.

:func:`~torch.abs`
:func:`~torch.asin`
:func:`~torch.asinh`
:func:`~torch.atan`
:func:`~torch.atanh`
:func:`~torch.ceil`
:func:`~torch.conj_physical`
:func:`~torch.floor`
:func:`~torch.log1p`
:func:`~torch.neg`
:func:`~torch.round`
:func:`~torch.sin`
:func:`~torch.sinh`
:func:`~torch.sign`
:func:`~torch.sgn`
:func:`~torch.signbit`
:func:`~torch.tan`
:func:`~torch.tanh`
:func:`~torch.trunc`
:func:`~torch.expm1`
:func:`~torch.sqrt`
:func:`~torch.angle`
:func:`~torch.isinf`
:func:`~torch.isposinf`
:func:`~torch.isneginf`
:func:`~torch.isnan`
:func:`~torch.erf`
:func:`~torch.erfinv`