# torch-sla: Differentiable Sparse Linear Algebra with Adjoint Solvers and Sparse Tensor Parallelism for PyTorch

arXiv: https://arxiv.org/abs/2601.13994 · HTML: https://mlsanity.com/paper/2601.13994/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 66 references that are in the map.
- First posted on arXiv 20 Jan 2026.

## Builds on (most cited first, 24 of 24)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Machine learning accelerated computational fluid dynamics](https://mlsanity.com/paper/2102.01010/) (2021)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [Message Passing Neural PDE Solvers](https://mlsanity.com/paper/2202.03376/) (2022)
- [DiffTaichi: Differentiable Programming for Physical Simulation](https://mlsanity.com/paper/1910.00935/) (2019)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [JFB: Jacobian-Free Backpropagation for Implicit Networks](https://mlsanity.com/paper/2103.12803/) (2021)
- [Learning to Optimize Multigrid PDE Solvers](https://mlsanity.com/paper/1902.10248/) (2019)
- [Theseus: A Library for Differentiable Nonlinear Optimization](https://mlsanity.com/paper/2207.09442/) (2022)
- [Geometry Aware Operator Transformer as an Efficient and Accurate Neural Surrogate for PDEs on Arbitrary Domains](https://mlsanity.com/paper/2505.18781/) (2025)
- [Learning Algebraic Multigrid Using Graph Neural Networks](https://mlsanity.com/paper/2003.05744/) (2020)
- [Optimistix: modular optimisation in JAX and Equinox](https://mlsanity.com/paper/2402.09983/) (2024)
- [Graph Neural Preconditioners for Iterative Solutions of Sparse Linear Systems](https://mlsanity.com/paper/2406.00809/) (2024)
- [Learning Interface Conditions in Domain Decomposition Solvers](https://mlsanity.com/paper/2205.09833/) (2022)
- [Mesh-Informed Neural Operator : A Transformer Generative Approach](https://mlsanity.com/paper/2506.16656/) (2025)

## Built on by (most cited first, 3 of 3)

- [Learning, Solving and Optimizing PDEs with TensorGalerkin: an efficient high-performance Galerkin assembly algorithm](https://mlsanity.com/paper/2602.05052/) (2026)
- [Preconditioned Physics-Informed Neural Operator Training](https://mlsanity.com/paper/2609.36216/) (2026)
- [On the Importance of Geometric Nonlinearity and Temperature-Dependent Properties in Multi-Material Thermo-Mechanical Topology Optimization](https://mlsanity.com/paper/2608.10344/) (2026)

## Abstract (arXiv)

Differentiable sparse linear algebra is foundational for scientific machine learning, yet PyTorch lacks a unified library for it: torch.sparse provides only low-level kernels and a non-differentiable, CPU-only spsolve, and torch.linalg is dense-only. We present torch-sla, an open-source library that fills this gap. It exposes a single autograd-aware API for direct, iterative, nonlinear, and eigenvalue solvers across five interchangeable backends -- SciPy and Eigen on CPU, cuDSS, CuPy, and a PyTorch-native iterative solver on GPU -- with automatic dispatch by device and problem size. The library further supports batched solves over shared or distinct sparsity patterns and distributed multi-GPU execution via domain decomposition with halo exchange. These capabilities are made scalable by an O(1)-graph adjoint differentiation framework and an autograd-compatible distributed halo-exchange layer. The library is available at https://www.torchsla.com/.