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

TL;DR

torch-sla is an open-source library that provides differentiable sparse linear algebra operations in PyTorch, supporting multiple backends, batched solves, and distributed multi-GPU execution.

Contribution

It introduces a unified, autograd-aware API for sparse linear algebra in PyTorch, with scalable differentiation and distributed computing capabilities.

Findings

Supports five interchangeable backends including SciPy, Eigen, cuDSS, CuPy, and PyTorch-native solvers.

Enables batched solves over shared or distinct sparsity patterns.

Provides scalable differentiation via an O(1)-graph adjoint framework and distributed halo exchange.

Abstract

Differentiable sparse linear algebra is foundational for scientific machine learning, yet PyTorch lacks a unified library for it: \texttt{torch.sparse} provides only low-level kernels and a non-differentiable, CPU-only \texttt{spsolve}, and \texttt{torch.linalg} is dense-only. We present \torchsla{}, 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/.