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.

Note

The specified elements and metadata mask of a semi-structured sparse tensor are stored together in a singleflat 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 firstm * k // 2 elements to be the kept elementsand the rest of the tensor is metadata.

In order to make it easier for the user to view the specified elementsand 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 typetorch.int16 if dtype is torch.float16 or torch.bfloat16, and element typetorch.int32 if dtype is torch.int8.

Note

It’s important to note thattorch.float32 is only supported for 1:2 sparsity. Therefore, it does not follow the same formula as above.

Constructing Sparse Semi-Structured Tensors#

You can transform a dense tensor into a sparse semi-structured tensor by simply using thetorch.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.

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 callto_sparse_semi_structured function to compress it for accelerated inference.

>>>fromtorch.sparseimportto_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 ofto_sparse_semi_structured(tensor) instead of usingtensor 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)))

Note

The memory consumption of a sparse COO tensor is at least(ndim*8+<sizeofelementtypeinbytes>)*nse bytes (plus a constantoverhead from storing other tensor data).

The memory consumption of a strided tensor is at leastproduct(<tensorshape>)*<sizeofelementtypeinbytes>.

For example, the memory consumption of a 10 000 x 10 000 tensorwith 100 000 non-zero 32-bit floating point numbers is at least(2*8+4)*100000=2000000 bytes when using COO tensorlayout and10000*10000*4=400000000 bytes when usingthe default strided tensor layout. Notice the 200 fold memorysaving from using the COO storage format.

Construction#

A sparse COO tensor can be constructed by providing the two tensors ofindices and values, as well as the size of the sparse tensor (when itcannot be inferred from the indices and values tensors) to a functiontorch.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))>>>stensor(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 inputi is NOT a list of index tuples. If you wantto write your indices this way, you should transpose before passing them tothe 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 sizeonly:

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

PyTorch implements an extension of sparse tensors with scalar valuesto sparse tensors with (contiguous) tensor values. Such tensors arecalled hybrid tensors.

PyTorch hybrid COO tensor extends the sparse COO tensor by allowingthevalues tensor to be a multi-dimensional tensor so that wehave:

• the indices of specified elements are collected inindicestensor of size(sparse_dims,nse) and with element typetorch.int64,

• the corresponding (tensor) values are collected invaluestensor of size(nse,dense_dims) and with an arbitrary integeror floating point number element type.

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))>>>stensor(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, ifs is a sparse COO tensor andM=s.sparse_dim(),K=s.dense_dim(), then we have the followinginvariants:

• M+K==len(s.shape)==s.ndim - dimensionality of a tensoris the sum of the number of sparse and dense dimensions,

• s.indices().shape==(M,nse) - sparse indices are storedexplicitly,

• s.values().shape==(nse,)+s.shape[M:M+K] - the valuesof a hybrid tensor are K-dimensional tensors,

• s.values().layout==torch.strided - values are stored asstrided tensors.

Uncoalesced sparse COO tensors#

PyTorch sparse COO tensor format permits sparseuncoalesced 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 allduplicate value entries. For example, one can specify multiple values,3 and4, for the same index1, that leads to an 1-Duncoalesced tensor:

>>>i=[[1,1]]>>>v=[3,4]>>>s=torch.sparse_coo_tensor(i,v,(3,))>>>stensor(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 elementsinto 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 oftorch.Tensor.coalesce() method is asparse tensor with the following properties:

• the indices of specified tensor elements are unique,

• the indices are sorted in lexicographical order,

• torch.Tensor.is_coalesced() returnsTrue.

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

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 atorch.Tensorinstance and to distinguish it from theTensor instances that usesome other layout, one can usetorch.Tensor.is_sparse ortorch.Tensor.layout properties:

>>>isinstance(s,torch.Tensor)True>>>s.is_sparseTrue>>>s.layout==torch.sparse_cooTrue

The number of sparse and dense dimensions can be acquired usingmethodstorch.Tensor.sparse_dim() andtorch.Tensor.dense_dim(), respectively. For instance:

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

Ifs is a sparse COO tensor then its COO format data can beacquired using methodstorch.Tensor.indices() andtorch.Tensor.values().

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

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

Constructing a new sparse COO tensor results a tensor that is notcoalesced:

>>>s.is_coalesced()False

but one can construct a coalesced copy of a sparse COO tensor usingthetorch.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 intoan account the additive nature of uncoalesced data: the values of thesame indices are the terms of a sum that evaluation gives the value ofthe corresponding tensor element. For example, the scalarmultiplication on a sparse uncoalesced tensor could be implemented bymultiplying all the uncoalesced values with the scalar becausec*(a+b)==c*a+c*b holds. However, any nonlinear operation,say, a square root, cannot be implemented by applying the operation touncoalesced data becausesqrt(a+b)==sqrt(a)+sqrt(b) does nothold in general.

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

>>>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 specifiedexplicitly and is assumed to be zero in general. However, there existsoperations that may interpret the fill value differently. Forinstance,torch.sparse.softmax() computes the softmax with theassumption that the fill value is negative infinity.

Note

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

Note

We say that an indices tensorcompressed_indices uses CSRencoding if the following invariants are satisfied:

• compressed_indices is a contiguous strided 32 or 64 bitinteger tensor

• compressed_indices shape is(*batchsize,compressed_dim_size+1) wherecompressed_dim_size is thenumber of compressed dimensions (e.g. rows or columns)

• compressed_indices[...,0]==0 where... denotes batchindices

• compressed_indices[...,compressed_dim_size]==nse wherense is the number of specified elements

• 0<=compressed_indices[...,i]-compressed_indices[...,i-1]<=plain_dim_size fori=1,...,compressed_dim_size,whereplain_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 tensorcreation viacheck_invariants=True keyword argument, orglobally usingtorch.sparse.check_sparse_tensor_invariantscontext manager instance. By default, the sparse tensor invariantschecks are disabled.

Note

The generalization of sparse compressed layouts to N-dimensionaltensors can lead to some confusion regarding the count of specifiedelements. When a sparse compressed tensor contains batch dimensionsthe number of specified elements will correspond to the number of suchelements per-batch. When a sparse compressed tensor has dense dimensionsthe element considered is now the K-dimensional array. Also for blocksparse compressed layouts the 2-D block is considered as the elementbeing specified. Take as an example a 3-dimensional block sparsetensor, with one batch dimension of lengthb, and a blockshape ofp,q. If this tensor hasn specified elements, thenin fact we haven blocks specified per batch. This tensor wouldhavevalues with shape(b,n,p,q). This interpretation of thenumber of specified elements comes from all sparse compressed layoutsbeing derived from the compression of a 2-dimensional matrix. Batchdimensions are treated as stacking of sparse matrices, dense dimensionschange the meaning of the element from a simple scalar value to anarray with its own dimensions.

Sparse CSR Tensor#

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

In the simplest case, a (0 + 2 + 0)-dimensional sparse CSR tensorconsists of three 1-D tensors:crow_indices,col_indices andvalues:

• Thecrow_indices tensor consists of compressed rowindices. This is a 1-D tensor of sizenrows+1 (the number ofrows plus 1). The last element ofcrow_indices is the numberof specified elements,nse. This tensor encodes the index invalues andcol_indices depending on where the given rowstarts. Each successive number in the tensor subtracted by thenumber before it denotes the number of elements in a given row.

• Thecol_indices tensor contains the column indices of eachelement. This is a 1-D tensor of sizense.

• Thevalues tensor contains the values of the CSR tensorelements. This is a 1-D tensor of sizense.

In the general case, the (B + 2 + K)-dimensional sparse CSR tensorconsists of two (B + 1)-dimensional index tensorscrow_indices andcol_indices, and of (1 + K)-dimensionalvalues tensor suchthat

• 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) wherelen(batchsize)==B andlen(densesize)==K.

>>>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)>>>csrtensor(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)

>>>a=torch.tensor([[0,0,1,0],[1,2,0,0],[0,0,0,0]],dtype=torch.float64)>>>sp=a.to_sparse_csr()>>>sptensor(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)

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

Sparse CSC Tensor#

The sparse CSC (Compressed Sparse Column) tensor format implements theCSC format for storage of 2 dimensional tensors with an extension tosupporting batches of sparse CSC tensors and values beingmulti-dimensional tensors.

Similarly tosparse CSR tensors, a sparse CSCtensor consists of three tensors:ccol_indices,row_indicesandvalues:

• Theccol_indices tensor consists of compressed columnindices. This is a (B + 1)-D tensor of shape(*batchsize,ncols+1).The last element is the number of specifiedelements,nse. This tensor encodes the index invalues androw_indices depending on where the given column starts. Eachsuccessive number in the tensor subtracted by the number before itdenotes the number of elements in a given column.

• Therow_indices tensor contains the row indices of eachelement. This is a (B + 1)-D tensor of shape(*batchsize,nse).

• Thevalues tensor contains the values of the CSC tensorelements. This is a (1 + K)-D tensor of shape(nse,*densesize).

>>>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)>>>csctensor(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)

>>>a=torch.tensor([[0,0,1,0],[1,2,0,0],[0,0,0,0]],dtype=torch.float64)>>>sp=a.to_sparse_csc()>>>sptensor(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 Tensor#

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

A sparse BSR tensor consists of three tensors:crow_indices,col_indices andvalues:

• Thecrow_indices tensor consists of compressed rowindices. 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 invalues andcol_indices depending on where the given column blockstarts. Each successive number in the tensor subtracted by thenumber before it denotes the number of blocks in a given row.

• Thecol_indices tensor contains the column block indices of eachelement. This is a (B + 1)-D tensor of shape(*batchsize,nse).

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

>>>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)>>>bsrtensor(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)

>>>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))>>>bsrtensor(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 Tensor#

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

A sparse BSC tensor consists of three tensors:ccol_indices,row_indices andvalues:

• Theccol_indices tensor consists of compressed columnindices. 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 invalues androw_indices depending on where the given row blockstarts. Each successive number in the tensor subtracted by thenumber before it denotes the number of blocks in a given column.

• Therow_indices tensor contains the row block indices of eachelement. This is a (B + 1)-D tensor of shape(*batchsize,nse).

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

>>>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)>>>bsctensor(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 conceptually very similar in that their indices data is splitinto two parts: so-called compressed indices that use the CSRencoding, and so-called plain indices that are orthogonal to thecompressed indices. This allows various tools on these tensors toshare the same implementations that are parameterized by tensorlayout.

>>>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)>>>csrtensor(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)>>>csctensor(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)

Linear Algebra operations#

The following table summarizes supported Linear Algebra operations onsparse matrices where the operands layouts may vary. HereT[layout] denotes a tensor with a given layout. Similarly,M[layout] denotes a matrix (2-D PyTorch tensor), andV[layout]denotes a vector (1-D PyTorch tensor). In addition,f denotes ascalar (float or 0-D PyTorch tensor),* is element-wisemultiplication, and@ is matrix multiplication.

where “Sparse grad?” column indicates if the PyTorch operation supportsbackward with respect to sparse matrix argument. All PyTorch operations,excepttorch.smm(), support backward with respect to stridedmatrix arguments.

Tensor methods and sparse#

The following Tensor methods are related to sparse tensors:

The following Tensor methods are specific to sparse COO tensors:

The following methods are specific tosparse CSR tensors andsparse BSR tensors:

The following methods are specific tosparse CSC tensors andsparse BSC tensors:

The following Tensor methods support sparse COO tensors:

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

Torch functions specific to sparse Tensors#

Other functions#

The followingtorch functions support sparse tensors:

cat()dstack()empty()empty_like()hstack()index_select()is_complex()is_floating_point()is_nonzero()is_same_size()is_signed()is_tensor()lobpcg()mm()native_norm()pca_lowrank()select()stack()svd_lowrank()unsqueeze()vstack()zeros()zeros_like()

To manage checking sparse tensor invariants, see:

To use sparse tensors withgradcheck() function,see:

Zero-preserving unary functions#

We aim to support all ‘zero-preserving unary functions’: functions of one argument that map zero to zero.

If you find that we are missing a zero-preserving unary functionthat 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.

abs()asin()asinh()atan()atanh()ceil()conj_physical()floor()log1p()neg()round()sin()sinh()sign()sgn()signbit()tan()tanh()trunc()expm1()sqrt()angle()isinf()isposinf()isneginf()isnan()erf()erfinv()