Skip to main content

​
Sparse Tensors

The torch.sparse module provides support for sparse tensors in multiple formats including COO, CSR, CSC, BSR, and BSC layouts.

​
Sparse Tensor Formats

PyTorch supports several sparse tensor formats:
  • COO (Coordinate format): Stores indices and values separately
  • CSR (Compressed Sparse Row): Row-compressed format
  • CSC (Compressed Sparse Column): Column-compressed format
  • BSR (Block Sparse Row): Block-based row-compressed format
  • BSC (Block Sparse Column): Block-based column-compressed format

​
Construction

​
Creating COO Tensors

torch.sparse_coo_tensor(
    indices,
    values,
    size=None,
    dtype=None,
    device=None,
    requires_grad=False,
    check_invariants=None
)
Constructs a sparse COO tensor.
​
indices
Tensor
2D tensor of shape (ndim, nnz) containing indices of non-zero elements.
​
values
Tensor
1D tensor of shape (nnz,) containing values.
​
size
tuple
default:"None"
Size of the sparse tensor. If None, inferred from indices.

​
Creating CSR Tensors

torch.sparse_csr_tensor(
    crow_indices,
    col_indices,
    values,
    size=None,
    dtype=None,
    device=None,
    requires_grad=False,
    check_invariants=None
)
Constructs a sparse CSR tensor.
​
crow_indices
Tensor
Compressed row indices.
​
col_indices
Tensor
Column indices.
​
values
Tensor
Values of non-zero elements.

​
Conversion

​
to_sparse

Tensor.to_sparse(
    sparseDims=None,
    *,
    layout=None,
    blocksize=None,
    dense_dim=None
)
Converts a dense tensor to sparse format.
​
layout
torch.layout
default:"None"
Desired sparse layout. Options: torch.sparse_coo, torch.sparse_csr, etc.
​
blocksize
tuple
default:"None"
Block size for block sparse formats (BSR/BSC).

​
to_dense

Tensor.to_dense(masked_grad=False)
Converts a sparse tensor to dense format.
​
masked_grad
bool
default:"False"
Whether to mask gradients by the sparsity pattern.

​
Arithmetic Operations

​
addmm

torch.sparse.addmm(
    mat,
    mat1,
    mat2,
    *,
    beta=1.0,
    alpha=1.0
)
Performs matrix multiplication with sparse matrix support: mat + alpha * (mat1 @ mat2) * beta.
​
mat
Tensor
Dense matrix to be added.
​
mat1
Tensor
Sparse matrix (COO or CSR) to be multiplied.
​
mat2
Tensor
Dense matrix to be multiplied.
​
beta
float
default:"1.0"
Multiplier for mat.
​
alpha
float
default:"1.0"
Multiplier for the matrix product.

​
mm

torch.sparse.mm(
    mat1,
    mat2,
    reduce='sum'
)
Performs matrix multiplication of sparse and dense/sparse matrices.
​
mat1
Tensor
Sparse matrix (first operand).
​
mat2
Tensor
Dense or sparse matrix (second operand).
​
reduce
str
default:"'sum'"
Reduction operation for duplicate indices. Options: 'sum', 'mean', 'amax', 'amin'.

​
sampled_addmm

torch.sparse.sampled_addmm(
    input,
    mat1,
    mat2,
    *,
    beta=1.0,
    alpha=1.0,
    out=None
)
Performs matrix multiplication sampled at the sparsity pattern of input.
​
input
Tensor
Sparse CSR tensor defining the sampling pattern.
​
mat1
Tensor
Dense matrix of shape (m, k).
​
mat2
Tensor
Dense matrix of shape (k, n).

​
Reduction Operations

​
sum

torch.sparse.sum(
    input,
    dim=None,
    dtype=None
)
Returns the sum of each row of the sparse tensor.
​
input
Tensor
Sparse tensor to sum.
​
dim
int or tuple
default:"None"
Dimension(s) to reduce. If None, reduce all dimensions.
​
dtype
torch.dtype
default:"None"
Desired data type of output.

​
softmax

torch.sparse.softmax(
    input,
    dim,
    *,
    dtype=None
)
Applies softmax to a sparse tensor along a dimension.
​
input
Tensor
Sparse tensor input.
​
dim
int
Dimension along which to apply softmax.

​
log_softmax

torch.sparse.log_softmax(
    input,
    dim,
    *,
    dtype=None
)
Applies log(softmax) to a sparse tensor.

​
Linear Algebra

​
spsolve (Sparse Solver)

torch.sparse.spsolve(
    input,
    other,
    *,
    left=True
)
Solves a sparse linear system.
​
input
Tensor
Sparse CSR coefficient matrix of shape (n, n).
​
other
Tensor
Dense right-hand side of shape (n,).
​
left
bool
default:"True"
Whether to solve input @ out = other (True) or out @ input = other (False).

​
Utility Functions

​
spdiags

torch.sparse.spdiags(
    diagonals,
    offsets,
    shape,
    layout=None
)
Creates a sparse 2D tensor by placing values from rows of diagonals along specified diagonals.
​
diagonals
Tensor
Matrix storing diagonals row-wise.
​
offsets
Tensor
The diagonals to be set (0 = main diagonal, negative = below, positive = above).
​
shape
tuple
Shape of the output tensor.
​
layout
torch.layout
default:"None"
Desired layout. Options: torch.sparse_coo, torch.sparse_csr, torch.sparse_csc.

​
Example Usage

import torch

# COO format
indices = torch.tensor([[0, 1, 2], [0, 1, 2]])
values = torch.tensor([1.0, 2.0, 3.0])
sparse_coo = torch.sparse_coo_tensor(indices, values, size=(3, 3))
print(f"COO sparse tensor:\n{sparse_coo}")
print(f"Dense representation:\n{sparse_coo.to_dense()}")

# CSR format
crow_indices = torch.tensor([0, 1, 2, 3])
col_indices = torch.tensor([0, 1, 2])
values = torch.tensor([1.0, 2.0, 3.0])
sparse_csr = torch.sparse_csr_tensor(crow_indices, col_indices, values, size=(3, 3))
print(f"\nCSR sparse tensor:\n{sparse_csr}")

# Convert dense to sparse
dense = torch.tensor([[1, 0, 0], [0, 2, 0], [0, 0, 3]])
sparse_from_dense = dense.to_sparse()
print(f"\nConverted to sparse:\n{sparse_from_dense}")

​
Best Practices

  • COO: Best for construction and random access patterns
  • CSR: Optimized for row-wise operations and matrix-vector products
  • CSC: Optimized for column-wise operations
  • BSR/BSC: Use for block-structured sparsity (e.g., in neural networks)
  • Convert to CSR/CSC for repeated arithmetic operations
  • Coalesce COO tensors before operations with .coalesce()
  • Use to_sparse() only when sparsity is significant (>90% zeros)
  • Batch sparse operations when possible
  • Consider dense operations for small tensors or low sparsity
  • COO @ Dense: Backward supported for both inputs
  • CSR @ Dense: Backward supported for both inputs
  • Sparse @ Sparse: Forward works, backward not supported
  • Use masked_grad=True in to_dense() to preserve sparsity pattern in gradients