mlsanityData to 13 Oct 2026

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

arXiv 2601.13994 ↗First posted 20 Jan 2026cs.DC cs.AIcited by 3

Not on the map. Why?

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/.

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

Builds on

24 of 66 references.

References · 24

  1. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,223
  2. 2016Semi-Supervised Classification with Graph Convolutional Networkscited by 11,318
  3. 2017Graph Attention Networkscited by 8,888
  4. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,637
  5. 2020Fourier Neural Operator for Parametric Partial Differential Equationscited by 2,857
  6. 2019Fast Graph Representation Learning with PyTorch Geometriccited by 2,422
  7. 2019DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operatorscited by 1,989
  8. 2020Learning Mesh-Based Simulation with Graph Networkscited by 852
  9. 2019Deep Equilibrium Modelscited by 606
  10. 2017OptNet: Differentiable Optimization as a Layer in Neural Networkscited by 592
  11. 2021Machine learning accelerated computational fluid dynamicscited by 465
  12. 2019Differentiable Convex Optimization Layerscited by 461
12 more
  1. 2022Message Passing Neural PDE Solverscited by 378
  2. 2019DiffTaichi: Differentiable Programming for Physical Simulationcited by 252
  3. 2021Efficient and Modular Implicit Differentiationcited by 201
  4. 2021JFB: Jacobian-Free Backpropagation for Implicit Networkscited by 86
  5. 2019Learning to Optimize Multigrid PDE Solverscited by 79
  6. 2022Theseus: A Library for Differentiable Nonlinear Optimizationcited by 64
  7. 2025Geometry Aware Operator Transformer as an Efficient and Accurate Neural Surrogate for PDEs on Arbitrary Domainscited by 59
  8. 2020Learning Algebraic Multigrid Using Graph Neural Networkscited by 54
  9. 2024Optimistix: modular optimisation in JAX and Equinoxcited by 18
  10. 2024Graph Neural Preconditioners for Iterative Solutions of Sparse Linear Systemscited by 15
  11. 2022Learning Interface Conditions in Domain Decomposition Solverscited by 11
  12. 2025Mesh-Informed Neural Operator : A Transformer Generative Approachcited by 11

Built on by

Cited by year (2026 to 13 Oct)

Citing papers per year: 2026: 3.

3 papers.

  1. 2026Learning, Solving and Optimizing PDEs with TensorGalerkin: an efficient high-performance Galerkin assembly algorithmcited by 3
  2. 2026Preconditioned Physics-Informed Neural Operator Trainingcited by 1
  3. 2026On the Importance of Geometric Nonlinearity and Temperature-Dependent Properties in Multi-Material Thermo-Mechanical Topology Optimizationcited by 0