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.

Performance Tips¶

Practical guidance for picking a backend, method, and precision once you are past the quick start. The reasoning behind the direct-vs-iterative trade-off is in Backends and Capability Matrix; the measured numbers are in Benchmarks.

Choosing a backend and method¶

float64 converges more reliably with the iterative solvers; float32 can stall on ill-conditioned PDE matrices.
For SPD matrices, cholesky is roughly twice as fast as lu – let matrix_type="auto" detect symmetry / positive-definiteness and pick it.
On CPU, scipy+lu is the default and gives machine precision.
On NVIDIA for problems under ~2M DOF, cudss+cholesky is the fastest direct solver.
For larger problems, pytorch+cg is the memory-efficient choice and the one that reached 169M DOF on a single GPU.
Beyond a single GPU, DSparseTensor partitions the matrix across devices.
For a portable GPU direct solve – including AMD ROCm, where cuDSS is not available – use strumpack (multifrontal LU on CPU/CUDA/ROCm).

Reusing work across solves¶

For repeated solves with the same matrix, factor once with A.lu() and reuse the LUFactorization for each right-hand side.
To apply one set of defaults across many solves, wrap them in a SolverConfig scope (see the Configuring solves section of Introduction) rather than repeating kwargs.
An iterative solve refines toward a tolerance, so a good x0 (e.g. the previous step’s solution in a time loop) cuts the iteration count.

Accuracy vs. speed¶

The iterative default atol ~ 1e-6 is a stopping point, not a ceiling – tighten atol toward 1e-12 when you need more digits, at the cost of iterations. See Direct vs iterative.
A good preconditioner ("amg" for PDE problems) can cut the iteration count 10-100x on ill-conditioned systems.