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.

torch-sla: PyTorch Sparse Linear Algebra¶

torch-sla solves sparse linear systems \(Ax = b\) in PyTorch. The matrix A is stored in sparse form, the solve runs on CPU or GPU, and gradients flow back through it via torch.autograd. It targets workloads that already live in PyTorch – FEM/CFD discretizations, physics-informed networks, graph operators – where you would otherwise copy out to SciPy or PETSc and lose the gradient.

The two core classes¶

Everything in torch-sla hangs off two tensor types. Both expose the same operation vocabulary (solve, eigsh, matvec, det, …); they differ only in where the matrix lives.

SparseTensor: A single-process sparse matrix in COO form, with an optional batch dimension. Runs on CPU or one GPU, dispatches to SciPy / PyTorch-native / cuDSS / STRUMPACK backends, and is differentiable through every solve. See SparseTensor.
DSparseTensor: A row-partitioned sparse matrix sharded across processes/GPUs with domain decomposition and halo exchange. Mirrors torch.distributed.tensor.DTensor and reuses the single-process operations rank-locally. See DSparseTensor.

For the per-operation reference – signatures, examples, sparsity figures, scaling plots – see Operations.

Quick start¶

pip install torch-sla

import torch
from torch_sla import SparseTensor

dense = torch.tensor([[ 4.0, -1.0,  0.0],
                      [-1.0,  4.0, -1.0],
                      [ 0.0, -1.0,  4.0]], dtype=torch.float64)
A = SparseTensor.from_dense(dense)

b = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64)
x = A.solve(b)               # auto-selects scipy+lu on CPU

Move the tensors to the GPU and the solve follows, picking cuDSS on NVIDIA:

x = A.cuda().solve(b.cuda())

To take gradients, set requires_grad on the values (or use the functional spsolve()) and call backward():

from torch_sla import spsolve

val = torch.tensor([...], requires_grad=True)
x = spsolve(val, row, col, shape, b)
x.sum().backward()           # grads w.r.t. val and b

How torch-sla compares¶

The tables below summarize where torch-sla fits relative to common alternatives. The short version: it earns its place when you need sparse solves inside a PyTorch autograd graph. If you are not using PyTorch, or you need an established preconditioner stack or true exascale distribution, the mature libraries below remain the better tools.

vs. `scipy.sparse.linalg`¶

Feature	torch-sla	scipy.sparse.linalg
PyTorch tensors	Native	Requires NumPy copy
GPU	cuDSS (NVIDIA), STRUMPACK (CUDA/ROCm)	CPU only
Gradients	Adjoint through solve/eig/svd	None
Batched solve	One call	Loop
Large scale	169M DOF measured (CG, single GPU)	Memory-bound on CPU
Distributed	DSparseTensor	No

vs. `torch.linalg.solve`¶

Feature	torch-sla	torch.linalg.solve
Matrix type	Sparse (COO/CSR)	Dense only
Memory (1M×1M, 1% density)	~80 MB	~8 TB (infeasible)
Solvers	LU, Cholesky, LDLT, CG, BiCGStab, GMRES, MINRES	Dense LU
Batching	Same or different patterns	Same shape only
Gradients	O(1) graph nodes via adjoint	Standard autograd

vs. NVIDIA AmgX¶

Feature	torch-sla	NVIDIA AmgX
Install	`pip install torch-sla`	Build from source
PyTorch integration	Native	Needs a wrapper
Gradients	Yes	No
Algebraic multigrid	Via the `amgx` backend (wraps AmgX) or `pyamg`	Core feature
Preconditioners	Jacobi, SSOR, polynomial, IC(0), AMG	ILU, AMG, and more

vs. PETSc¶

Feature	torch-sla	PETSc
Install	`pip install`	MPI + compilers
PyTorch integration	Native tensors	petsc4py + copies
Gradients	Yes	No
Solver coverage	Core Krylov + direct	Extensive (KSP, SNES)
Distributed	DSparseTensor, multi-GPU	Full MPI, exascale-proven

PETSc remains the right choice past the scales torch-sla has been tested at, or when you need its solver and preconditioner breadth.

Indices and search¶

License¶

Contact¶

Author: Mingyuan Chi
Email: walker.chi.000@gmail.com
Author: Shizheng Wen
Email: shizheng.wen@sam.math.ethz.ch

Citation¶

@article{chi2026torchsla,
  title={torch-sla: Differentiable Sparse Linear Algebra with Adjoint Solvers and Sparse Tensor Parallelism for PyTorch},
  author={Chi, Mingyuan and Wen, Shizheng},
  journal={arXiv preprint arXiv:2601.13994},
  year={2026},
  url={https://arxiv.org/abs/2601.13994}
}

arXiv:2601.13994 — Differentiable Sparse Linear Algebra with Adjoint Solvers and Sparse Tensor Parallelism for PyTorch.