|  | .. 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` elements contiguously in | 
|  | 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, semi-structured, 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>`. | 
|  |  | 
|  | We also have a prototype implementation to support :ref: `semi-structured sparsity<sparse-semi-structured-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-semi-structured-docs: | 
|  |  | 
|  | Sparse Semi-Structured Tensors | 
|  | ++++++++++++++++++++++++++++++ | 
|  |  | 
|  | .. warning:: | 
|  |  | 
|  | Sparse semi-structured tensors are currently a prototype feature and subject to change. Please feel free to open an issue to report a bug or if you have feedback to share. | 
|  |  | 
|  | Semi-Structured sparsity is a sparse data layout that was first introduced in NVIDIA's Ampere architecture. It is also referred to as **fine-grained structured sparsity** or **2:4 structured sparsity**. | 
|  |  | 
|  | This sparse layout stores `n` elements out of every `2n` elements, with `n` being determined by the width of the Tensor's data type (dtype). The most frequently used dtype is float16, where `n=2`, thus the term "2:4 structured sparsity." | 
|  |  | 
|  | Semi-structured sparsity is explained in greater detail in `this NVIDIA blog post <https://developer.nvidia.com/blog/exploiting-ampere-structured-sparsity-with-cusparselt>`_. | 
|  |  | 
|  | In PyTorch, semi-structured sparsity is implemented via a Tensor subclass. | 
|  | By subclassing, we can override ``__torch_dispatch__`` , allowing us to use faster sparse kernels when performing matrix multiplication. | 
|  | We can also store the tensor in it's compressed form inside the subclass to reduce memory overhead. | 
|  |  | 
|  | In this compressed form, the sparse tensor is stored by retaining only the *specified* elements and some metadata, which encodes the mask. | 
|  |  | 
|  | .. note:: | 
|  | The specified elements and metadata mask of a semi-structured sparse tensor are stored together in a single | 
|  | flat compressed tensor. They are appended to each other to form a contiguous chunk of memory. | 
|  |  | 
|  | compressed tensor = [ specified elements of original tensor |   metadata_mask ] | 
|  |  | 
|  | For an original tensor of size `(r, c)` we expect the first `m * k // 2` elements to be the kept elements | 
|  | and the rest of the tensor is metadata. | 
|  |  | 
|  | In order to make it easier for the user to view the specified elements | 
|  | and mask, one can use ``.indices()`` and ``.values()`` to access the mask and specified elements respectively. | 
|  |  | 
|  |  | 
|  | - ``.values()`` returns the specified elements in a tensor of size `(r, c//2)` and with the same dtype as the dense matrix. | 
|  |  | 
|  | - ``.indices()`` returns the metadata_mask in a tensor of size `(r, c//2 )` and with element type ``torch.int16`` if dtype is torch.float16 or torch.bfloat16, and element type ``torch.int32`` if dtype is torch.int8. | 
|  |  | 
|  |  | 
|  | For 2:4 sparse tensors, the metadata overhead is minor - just 2 bits per specified element. | 
|  |  | 
|  | .. note:: | 
|  | It's important to note that ``torch.float32`` is only supported for 1:2 sparsity. Therefore, it does not follow the same formula as above. | 
|  |  | 
|  | Here, we break down how to calculate the compression ratio ( size dense / size sparse) of a 2:4 sparse tensor. | 
|  |  | 
|  | Let `(r, c) = tensor.shape` and `e = bitwidth(tensor.dtype)`, so `e = 16` for ``torch.float16`` and ``torch.bfloat16`` and `e = 8` for ``torch.int8``. | 
|  |  | 
|  | .. math:: | 
|  | M_{dense} = r \times c \times e \\ | 
|  | M_{sparse} = M_{specified} + M_{metadata} = r \times \frac{c}{2} \times e + r \times \frac{c}{2} \times 2 = \frac{rce}{2} + rc =rce(\frac{1}{2} +\frac{1}{e}) | 
|  |  | 
|  | Using these calculations, we can determine the total memory footprint for both the original dense and the new sparse representation. | 
|  |  | 
|  | This gives us a simple formula for the compression ratio, which is dependent only on the bitwidth of the tensor datatype. | 
|  |  | 
|  | .. math:: | 
|  | C = \frac{M_{sparse}}{M_{dense}} =  \frac{1}{2} + \frac{1}{e} | 
|  |  | 
|  | By using this formula, we find that the compression ratio is 56.25% for ``torch.float16`` or ``torch.bfloat16``, and 62.5% for ``torch.int8``. | 
|  |  | 
|  | Constructing Sparse Semi-Structured Tensors | 
|  | ------------------------------------------- | 
|  |  | 
|  | You can transform a dense tensor into a sparse semi-structured tensor by simply using the ``torch.to_sparse_semi_structured`` function. | 
|  |  | 
|  | Please also note that we only support CUDA tensors since hardware compatibility for semi-structured sparsity is limited to NVIDIA GPUs. | 
|  |  | 
|  |  | 
|  | The following datatypes are supported for semi-structured sparsity. Note that each datatype has its own shape constraints and compression factor. | 
|  |  | 
|  | .. csv-table:: | 
|  | :header: "PyTorch dtype", "Shape Constraints", "Compression Factor", "Sparsity Pattern" | 
|  | :widths: 15, 45, 10, 10 | 
|  | :delim: ; | 
|  |  | 
|  | ``torch.float16``; Tensor must be 2D and (r, c) must both be a positive multiple of 64;9/16;2:4 | 
|  | ``torch.bfloat16``; Tensor must be 2D and (r, c) must both be a positive multiple of 64;9/16;2:4 | 
|  | ``torch.int8``; Tensor must be 2D and (r, c) must both be a positive multiple of 128;10/16;2:4 | 
|  |  | 
|  |  | 
|  | To construct a semi-structured sparse tensor, start by creating a regular dense tensor that adheres to a 2:4 (or semi-structured) sparse format. | 
|  | To do this we  tile a small 1x4 strip to create a 16x16 dense float16 tensor. | 
|  | Afterwards, we can call ``to_sparse_semi_structured`` function to compress it for accelerated inference. | 
|  |  | 
|  | >>> from torch.sparse import to_sparse_semi_structured | 
|  | >>> A = torch.Tensor([0, 0, 1, 1]).tile((128, 32)).half().cuda() | 
|  | tensor([[0., 0., 1.,  ..., 0., 1., 1.], | 
|  | [0., 0., 1.,  ..., 0., 1., 1.], | 
|  | [0., 0., 1.,  ..., 0., 1., 1.], | 
|  | ..., | 
|  | [0., 0., 1.,  ..., 0., 1., 1.], | 
|  | [0., 0., 1.,  ..., 0., 1., 1.], | 
|  | [0., 0., 1.,  ..., 0., 1., 1.]], device='cuda:0', dtype=torch.float16) | 
|  | >>> A_sparse = to_sparse_semi_structured(A) | 
|  | SparseSemiStructuredTensor(shape=torch.Size([128, 128]), transposed=False, values=tensor([[1., 1., 1.,  ..., 1., 1., 1.], | 
|  | [1., 1., 1.,  ..., 1., 1., 1.], | 
|  | [1., 1., 1.,  ..., 1., 1., 1.], | 
|  | ..., | 
|  | [1., 1., 1.,  ..., 1., 1., 1.], | 
|  | [1., 1., 1.,  ..., 1., 1., 1.], | 
|  | [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0', dtype=torch.float16), metadata=tensor([[-4370, -4370, -4370,  ..., -4370, -4370, -4370], | 
|  | [-4370, -4370, -4370,  ..., -4370, -4370, -4370], | 
|  | [-4370, -4370, -4370,  ..., -4370, -4370, -4370], | 
|  | ..., | 
|  | [-4370, -4370, -4370,  ..., -4370, -4370, -4370], | 
|  | [-4370, -4370, -4370,  ..., -4370, -4370, -4370], | 
|  | [-4370, -4370, -4370,  ..., -4370, -4370, -4370]], device='cuda:0', | 
|  | dtype=torch.int16)) | 
|  |  | 
|  | Sparse Semi-Structured Tensor Operations | 
|  | ---------------------------------------- | 
|  |  | 
|  | Currently, the following operations are supported for semi-structured sparse tensors: | 
|  |  | 
|  | - torch.addmm(bias, dense, sparse.t()) | 
|  | - torch.mm(dense, sparse) | 
|  | - torch.mm(sparse, dense) | 
|  | - aten.linear.default(dense, sparse, bias) | 
|  | - aten.t.default(sparse) | 
|  | - aten.t.detach(sparse) | 
|  |  | 
|  | To use these ops, simply pass the output of ``to_sparse_semi_structured(tensor)``  instead of using ``tensor`` once your tensor has 0s in a semi-structured sparse format, like this: | 
|  |  | 
|  | >>> a = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).half().cuda() | 
|  | >>> b = torch.rand(64, 64).half().cuda() | 
|  | >>> c = torch.mm(a, b) | 
|  | >>> a_sparse = to_sparse_semi_structured(a) | 
|  | >>> torch.allclose(c, torch.mm(a_sparse, b)) | 
|  | True | 
|  |  | 
|  | Accelerating nn.Linear with semi-structured sparsity | 
|  | ---------------------------------------------------- | 
|  | You can accelerate the linear layers in your model if the weights are already semi-structured sparse with just a few lines of code: | 
|  |  | 
|  | >>> input = torch.rand(64, 64).half().cuda() | 
|  | >>> mask = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).cuda().bool() | 
|  | >>> linear = nn.Linear(64, 64).half().cuda() | 
|  | >>> linear.weight = nn.Parameter(to_sparse_semi_structured(linear.weight.masked_fill(~mask, 0))) | 
|  |  | 
|  |  | 
|  | .. _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, one 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.matmul`; no; ``M[SparseSemiStructured] @ M[strided] -> M[strided]`` | 
|  | :func:`torch.matmul`; no; ``M[strided] @ M[SparseSemiStructured] -> M[strided]`` | 
|  | :func:`torch.mm`; no; ``M[sparse_coo] @ M[strided] -> M[strided]`` | 
|  | :func:`torch.mm`; no; ``M[SparseSemiStructured] @ M[strided] -> M[strided]`` | 
|  | :func:`torch.mm`; no; ``M[strided] @ M[SparseSemiStructured] -> 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.addmm`; no; ``f * M[strided] + f * (M[SparseSemiStructured] @ M[strided]) -> M[strided]`` | 
|  | :func:`torch.addmm`; no; ``f * M[strided] + f * (M[strided] @ M[SparseSemiStructured]) -> 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 | 
|  |  | 
|  | To use sparse tensors with :func:`~torch.autograd.gradcheck` function, | 
|  | see: | 
|  |  | 
|  | .. autosummary:: | 
|  | :toctree: generated | 
|  | :nosignatures: | 
|  |  | 
|  | sparse.as_sparse_gradcheck | 
|  |  | 
|  | 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` |