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) and Eigen (CG, BiCGStab). GPU solvers include CuPy (LU, CG, GMRES) 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 or CuPy 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.

Installation¶

PyPI: pip install torch-sla — simplest installation
GitHub: Clone and install for development
Optional backends: cuDSS, Eigen for enhanced performance

Using pip¶

To install the latest release:

pip install torch-sla

Or install from GitHub for the latest development version:

pip install git+https://github.com/walkerchi/torch-sla.git

Optional Dependencies¶

The core install pulls in torch, numpy, scipy, and ninja — enough to run CPU solvers out of the box. GPU users can pick the backend(s) they need:

# GPU users: choose one or both CUDA 12+ backends
pip install torch-sla[cupy]     # + CuPy backend
pip install torch-sla[cudss]    # + cuDSS backend (fastest direct solver on GPU)

# Full installation with all runtime backends (does not include dev/docs)
pip install torch-sla[all]

# For development tools (pytest, black, isort, mypy)
pip install torch-sla[dev]

# For documentation tools (sphinx, furo)
pip install torch-sla[docs]

Verify your environment

After installation, you can inspect which backends are available on your machine:

import torch_sla
torch_sla.show_backends()</code></pre>
</div>



Backend Requirements¶








Backend
Installation
Notes



scipy
pip install scipy
Default, always available

pytorch
Included with PyTorch
Native CG/BiCGStab solvers

cupy
pip install cupy-cuda12x
GPU direct + iterative solvers via cupyx.scipy

cudss
pip install nvmath-python[cu12]
Best for medium-scale GPU problems (10K-2M DOF)

Backend	Installation	Notes
`scipy`	`pip install scipy`	Default, always available
`pytorch`	Included with PyTorch	Native CG/BiCGStab solvers
`cupy`	`pip install cupy-cuda12x`	GPU direct + iterative solvers via cupyx.scipy
`cudss`	`pip install nvmath-python[cu12]`	Best for medium-scale GPU problems (10K-2M DOF)