torch-sla (Torch Sparse Linear Algebra) is a Python library that provides differentiable sparse linear equation solvers for PyTorch. It solves systems of the form Ax = b where A is a sparse matrix, with full support for automatic differentiation (autograd) and GPU acceleration via CUDA.

How do I solve a sparse linear system in PyTorch?

Use torch-sla's SparseTensor class: from torch_sla import SparseTensor; A = SparseTensor(values, row, col, shape); x = A.solve(b). This works on both CPU and GPU, and supports gradient computation.

What sparse solvers does torch-sla support?

torch-sla supports multiple backends: CPU solvers include SciPy (LU, UMFPACK, CG, BiCGStab, GMRES, MINRES) and PyTorch-native (CG, BiCGStab, GMRES, MINRES). GPU solvers include PyTorch-native (CG, BiCGStab, GMRES, MINRES) and cuDSS (LU, Cholesky, LDLT). The library automatically selects the best solver based on hardware and matrix properties.

Can I compute gradients through sparse solve in PyTorch?

Yes, torch-sla fully supports PyTorch autograd. You can set requires_grad=True on your sparse matrix values, solve the system, compute a loss, and call backward() to get gradients with respect to the matrix values and right-hand side.

How do I solve batched sparse systems in PyTorch?

torch-sla supports batched solving for matrices with the same sparsity pattern. Create a SparseTensor with batched values of shape [batch_size, nnz] and solve all systems in parallel. For matrices with different patterns, use SparseTensorList.

How do I use torch-sla on GPU?

Simply call .cuda() on your SparseTensor: A_cuda = A.cuda(); x = A_cuda.solve(b.cuda()). The library automatically uses cuDSS for GPU-accelerated solving.

What is the difference between torch-sla and scipy.sparse?

torch-sla offers native PyTorch integration, GPU support via CUDA, full autograd gradient support, and batched solving. scipy.sparse is CPU-only, requires data copying to/from PyTorch, and doesn't support automatic differentiation.

How do I install torch-sla?

Install via pip: pip install torch-sla. For GPU support, ensure you have CUDA installed and a compatible PyTorch version.

SparseTensor¶

SparseTensor is the single-process sparse matrix. It stores a matrix in COO form (val, row, col plus a shape), optionally with a leading batch dimension, and carries the full operation vocabulary: linear and nonlinear solves, eigendecomposition, SVD, matrix–vector products, scalar/structural queries, graph analysis and visualization. Every solve is differentiable through torch.autograd and dispatches to the backend best suited to the device (SciPy / PyTorch-native on CPU, cuDSS / STRUMPACK on GPU).

For the matrix sharded across processes, see DSparseTensor. For the per-operation reference, see Operations.

Construction¶

A SparseTensor can be built directly from COO triplets or through several convenience constructors:

import torch
from torch_sla import SparseTensor

# Direct COO: values, row indices, col indices, shape
A = SparseTensor(val, row, col, (n, n))

# From a dense matrix (drops the zeros)
A = SparseTensor.from_dense(dense)

# From an explicit list of (row, col, value) entries
A = SparseTensor.from_coo_list(entries, shape=(n, n))

# Structured constructors
I  = SparseTensor.eye(n)                 # identity
D  = SparseTensor.diag(values)           # diagonal
T  = SparseTensor.tridiagonal(n, ...)    # tridiagonal band

Conversions back out:

dense = A.to_dense()          # torch.Tensor
crow, col, val = A.to_csr()   # CSR triplet
t = A.to_torch_sparse()       # torch.sparse_coo_tensor

A leading batch dimension (shape [B, n, n]) makes every operation act on a stack of same-pattern matrices in one call; see solve_batch. For matrices with different patterns use SparseTensorList.

Operation catalog¶

Each operation links to its detailed entry in Operations and to its API object. Operations marked (grad) propagate gradients via the adjoint method.

Linear solves¶

Operation	API	Description
solve (grad)	`solve()`	Solve \(Ax = b\); auto-selects a direct or iterative backend.
solve_batch	`solve_batch()`	Many right-hand sides / value sets sharing one sparsity pattern.
lu	`lu()`	Cache an LU factorization for repeated solves.

Nonlinear¶

Operation	API	Description
nonlinear_solve (grad)	`nonlinear_solve()`	Newton / Picard / Anderson solve of \(F(u, \theta) = 0\).

Eigen / spectral¶

Operation	API	Description
eigsh (grad)	`eigsh()`	Top-k eigenpairs of a symmetric/Hermitian matrix.
eigs	`eigs()`	Top-k eigenpairs of a general matrix.
svd (grad)	`svd()`	Truncated rank-k singular value decomposition.

Matrix–vector¶

Operation	API	Description
matvec / @	`__matmul__()`	Sparse matrix–vector / matrix–matrix product (SpMV).

Scalar / structural¶

Operation	API	Description
det (grad)	`det()`	Determinant via sparse LU.
logdet (grad)	`logdet()`	Log-determinant (numerically stable for large matrices).
norm (grad)	`norm()`	Frobenius / 1- / 2-norm.
condition_number	`condition_number()`	Ratio \(\sigma_{\max}/\sigma_{\min}\).
is_symmetric / is_positive_definite	`is_symmetric()`	Structural predicates for solver selection.

Graph¶

Operation	API	Description
connected_components	`connected_components()`	Label connected components of the adjacency pattern.

Visualization¶

Operation	API	Description
spy	`spy()`	Plot the sparsity pattern as a matplotlib figure.

I/O and reductions¶

Beyond the headline operations, SparseTensor also offers save/load (safetensors and Matrix Market), element-wise math (abs, sqrt, exp, log, …) and reductions (sum, mean, max, min). These are documented in the API reference.