API Reference

This section provides the complete API documentation for torch-sla.

---

Core Classes

SparseTensor

The main class for working with sparse matrices. Supports batched operations, automatic differentiation, and multiple backends.

class torch_sla.SparseTensor(values: Tensor, row_indices: Tensor, col_indices: Tensor, shape: Tuple[int, ...], sparse_dim: Tuple[int, int] = (-2, -1))

Bases: object

Wrapper class for PyTorch sparse tensors with batched and block support.

Supports tensors with shape […batch, M, N, …block] where: - Leading dimensions […batch] are batch dimensions - (M, N) are the sparse matrix dimensions (at sparse_dim positions) - Trailing dimensions […block] are block dimensions

Parameters:

• values (torch.Tensor) – Non-zero values with shape: - Simple: [nnz] - Batched: […batch, nnz] - Block: [nnz, *block_shape] - Batched+Block: […batch, nnz, *block_shape]

• row_indices (torch.Tensor) – Row indices with shape [nnz]. Must be on the same device as values.

• col_indices (torch.Tensor) – Column indices with shape [nnz]. Must be on the same device as values.

• shape (Tuple[int, ...]) – Full tensor shape […batch, M, N, *block_shape].

• sparse_dim (Tuple[int, int], optional) – Which dimensions are sparse (M, N). Default: (-2, -1) meaning last two before any block dimensions.

values

The non-zero values.

Type:

torch.Tensor

row_indices

Row indices of non-zeros.

Type:

torch.Tensor

col_indices

Column indices of non-zeros.

Type:

torch.Tensor

shape

Full tensor shape.

Type:

Tuple[int, …]

sparse_shape

The (M, N) dimensions.

Type:

Tuple[int, int]

batch_shape

The batch dimensions.

Type:

Tuple[int, …]

block_shape

The block dimensions.

Type:

Tuple[int, …]

Examples

1. Simple 2D Sparse Matrix [M, N]

>>> import torch
>>> from torch_sla import SparseTensor
>>>
>>> # Create a 3x3 tridiagonal matrix in COO format
>>> val = torch.tensor([4.0, -1.0, -1.0, 4.0, -1.0, -1.0, 4.0])
>>> row = torch.tensor([0, 0, 1, 1, 1, 2, 2])
>>> col = torch.tensor([0, 1, 0, 1, 2, 1, 2])
>>> A = SparseTensor(val, row, col, (3, 3))
>>> print(A)
SparseTensor(shape=(3, 3), sparse=(3, 3), nnz=7, dtype=torch.float64, device=cpu)
>>>
>>> # Solve Ax = b
>>> b = torch.tensor([1.0, 2.0, 3.0])
>>> x = A.solve(b)

2. Batched Sparse Matrices [B, M, N]

Same sparsity pattern, different values for each batch.

>>> # 4 matrices, each 3x3, same structure
>>> batch_size = 4
>>> val_batch = val.unsqueeze(0).expand(batch_size, -1).clone()  # [4, 7]
>>> for i in range(batch_size):
...     val_batch[i] = val * (1.0 + 0.1 * i)  # Scale each matrix
>>>
>>> A_batch = SparseTensor(val_batch, row, col, (4, 3, 3))
>>> print(A_batch.batch_shape)  # (4,)
>>> print(A_batch.sparse_shape)  # (3, 3)
>>>
>>> # Batched solve
>>> b_batch = torch.randn(4, 3)
>>> x_batch = A_batch.solve(b_batch)  # [4, 3]

3. Multi-Dimensional Batch [B1, B2, M, N]

>>> B1, B2 = 2, 3  # e.g., 2 materials x 3 temperatures
>>> val_batch = val.unsqueeze(0).unsqueeze(0).expand(B1, B2, -1).clone()  # [2, 3, 7]
>>> A_multi = SparseTensor(val_batch, row, col, (B1, B2, 3, 3))
>>> print(A_multi.batch_shape)  # (2, 3)
>>>
>>> b_multi = torch.randn(B1, B2, 3)
>>> x_multi = A_multi.solve(b_multi)  # [2, 3, 3]

4. Block Sparse Matrix [M, N, K, K] (Block Size K)

Each non-zero entry is a KxK dense block instead of a scalar.

>>> # 2x2 block matrix with 2x2 blocks = 4x4 total
>>> block_size = 2
>>> nnz = 3  # 3 non-zero blocks
>>>
>>> # Values: [nnz, K, K] = [3, 2, 2]
>>> val_block = torch.randn(nnz, block_size, block_size)
>>> row_block = torch.tensor([0, 0, 1])  # Block row indices
>>> col_block = torch.tensor([0, 1, 1])  # Block col indices
>>>
>>> # Shape: (num_block_rows, num_block_cols, block_size, block_size)
>>> A_block = SparseTensor(val_block, row_block, col_block, (2, 2, 2, 2))
>>> print(A_block.block_shape)  # (2, 2)
>>> print(A_block.sparse_shape)  # (2, 2) - number of blocks
>>> print(A_block.shape)  # (2, 2, 2, 2) - full shape

5. Batched Block Sparse [B, M, N, K, K]

>>> batch_size = 4
>>> val_batch_block = torch.randn(batch_size, nnz, block_size, block_size)  # [4, 3, 2, 2]
>>> A_batch_block = SparseTensor(val_batch_block, row_block, col_block, (4, 2, 2, 2, 2))
>>> print(A_batch_block.batch_shape)  # (4,)
>>> print(A_batch_block.block_shape)  # (2, 2)

6. Create from Dense Matrix

>>> A_dense = torch.randn(100, 100)
>>> A_dense[A_dense.abs() < 0.5] = 0  # Sparsify
>>> A = SparseTensor.from_dense(A_dense)

7. Create from PyTorch Sparse Tensor

>>> A_torch = torch.randn(100, 100).to_sparse_coo()
>>> A = SparseTensor.from_torch_sparse(A_torch)

8. Property Detection

>>> A = SparseTensor(val, row, col, (3, 3))
>>> A.is_symmetric()  # tensor(True) - returns tensor for batch support
>>> A.is_positive_definite()  # tensor(True)
>>> A.is_positive_definite('cholesky')  # Use Cholesky factorization check

9. Matrix Operations

>>> # Matrix-vector multiply
>>> y = A @ x  # SparseTensor @ dense vector
>>>
>>> # Sparse-sparse multiply (returns SparseTensor with sparse gradients)
>>> C = A @ A
>>>
>>> # Norms
>>> A.norm('fro')  # Frobenius norm
>>>
>>> # Eigenvalues (symmetric matrices)
>>> eigenvalues, eigenvectors = A.eigsh(k=2, which='LM')

10. CUDA Support

>>> A_cuda = A.cuda()
>>> x = A_cuda.solve(b.cuda())  # Uses cuDSS

classmethod from_dense(A: Tensor, sparse_dim: Tuple[int, int] = (-2, -1)) → SparseTensor

Create SparseTensor from dense tensor.

Parameters:

• A (torch.Tensor) – Dense tensor with shape […batch, M, N, …block].

• sparse_dim (Tuple[int, int], optional) – Which dimensions are sparse. Default: (-2, -1).

Returns:

Sparse representation of A.

Return type:

SparseTensor

Examples

>>> A_dense = torch.randn(3, 3)
>>> A_dense[A_dense.abs() < 0.5] = 0
>>> A = SparseTensor.from_dense(A_dense)

classmethod from_torch_sparse(A: Tensor) → SparseTensor

Create SparseTensor from PyTorch sparse tensor.

Parameters:

A (torch.Tensor) – PyTorch sparse COO or CSR tensor (2D only).

Returns:

SparseTensor representation.

Return type:

SparseTensor

Examples

>>> A_coo = torch.randn(3, 3).to_sparse_coo()
>>> A = SparseTensor.from_torch_sparse(A_coo)

classmethod eye(n: int, dtype: dtype = torch.float64, device: str | device = 'cpu') → SparseTensor

Sparse identity n x n.

classmethod diag(values: Tensor, device: str | device | None = None) → SparseTensor

Sparse diagonal matrix from a 1-D vector.

classmethod tridiagonal(n: int, diag: float | Tensor = 2.0, off_diag: float | Tensor = -1.0, dtype: dtype = torch.float64, device: str | device = 'cpu') → SparseTensor

Sparse symmetric tridiagonal n x n. diag=4, off=-1 is the canonical SPD test matrix; diag=2, off=-1 is the 1-D Laplacian. diag / off_diag accept scalars or matching-length tensors.

property shape: Tuple[int, ...]

Full tensor shape […batch, M, N, …block].

property sparse_shape: Tuple[int, int]

The (M, N) sparse matrix dimensions.

property batch_shape: Tuple[int, ...]

The batch dimensions before the sparse dimensions.

property block_shape: Tuple[int, ...]

The block dimensions after the sparse dimensions.

property sparse_dim: Tuple[int, int]

The dimensions that are sparse (M, N).

property ndim: int

Number of dimensions.

property nnz: int

Number of non-zero elements (per batch/block).

property dtype: dtype

Data type of the values.

property device: device

Device of the tensor.

property is_cuda: bool

Whether the tensor is on CUDA.

property is_batched: bool

Whether the tensor has batch dimensions.

property is_block: bool

Whether the tensor has block dimensions.

property batch_size: int

Total number of batch elements (product of batch_shape).

property is_square: bool

Whether the sparse dimensions are square (M == N).

to(device: str | device | None = None, dtype: dtype | None = None) → SparseTensor

Move tensor to device and/or convert dtype.

Parameters:

• device (str or torch.device, optional) – Target device (e.g., ‘cuda’, ‘cpu’, ‘cuda:0’).

• dtype (torch.dtype, optional) – Target data type (e.g., torch.float32, torch.float64).

Returns:

New SparseTensor on the target device/dtype.

Return type:

SparseTensor

Examples

>>> A = SparseTensor(val, row, col, shape)
>>> A_cuda = A.to('cuda')
>>> A_float32 = A.to(dtype=torch.float32)
>>> A_cuda_float32 = A.to('cuda', torch.float32)

cuda(device: int | None = None) → SparseTensor

Move tensor to CUDA device.

Parameters:

device (int, optional) – CUDA device index. Default: current device.

Returns:

Tensor on CUDA.

Return type:

SparseTensor

cpu() → SparseTensor

Move tensor to CPU.

Returns:

Tensor on CPU.

Return type:

SparseTensor

float() → SparseTensor

Convert to float32.

double() → SparseTensor

Convert to float64.

half() → SparseTensor

Convert to float16.

requires_grad_(requires_grad: bool = True) → SparseTensor

Mark the underlying values tensor as requiring gradient.

Mirrors torch.Tensor.requires_grad_(): returns self so the call can be chained, and flips self.values.requires_grad in-place. Indices (row_indices / col_indices) are non-differentiable by construction and unaffected.

The 0.2.x line carried this as a bound method; the 0.3 refactor moved the body into ops but forgot to re-expose a delegating shim on the class, breaking every downstream caller that did A.requires_grad_(True) (TensorMesh backward tests in particular). Restore the delegation.

to_torch_sparse(*args, **kwargs)

to_dense(*args, **kwargs)

to_csr(*args, **kwargs)

extract_partition(*args, **kwargs)

save_distributed(*args, **kwargs)

partition_for_rank(*args, **kwargs)

detect_matrix_type(*args, **kwargs)

T(*args, **kwargs)

conj(*args, **kwargs)

H(*args, **kwargs)

flatten_blocks(*args, **kwargs)

unflatten_blocks(*args, **kwargs)

is_symmetric(*args, **kwargs)

is_hermitian(*args, **kwargs)

is_positive_definite(*args, **kwargs)

connected_components(*args, **kwargs)

Find connected components of the graph \(G(A)\) of this sparse matrix.

\[G(A) = (V, E),\quad V=\{0,\dots,N-1\},\quad E = \{(i,j) : A_{ij}\neq 0\}\]

Partition \(V\) into maximal subsets connected by E (treated as undirected). Returns a label per node and the component count.

Algorithm

FastSV / Shiloach-Vishkin label propagation with pointer jumping: each node starts as its own root; every round hooks each edge to the minimum neighbouring label, then path-compresses all the way to the root (so convergence is in graph depth-independent \(O(\log N)\) rounds), all vectorized on-device:

label[i] = i for all i
repeat until labels stable (checksum unchanged):
    # hooking: each edge pulls the smaller root label
    label[dst] = min(label[dst], label[src])   # scatter_reduce amin
    # pointer jumping: label[i] = label[label[i]] until rooted
    compress_to_roots(label)
relabel to contiguous 0..num_components-1

Complexity

Time \(O(nnz \log N)\) (O(log N) rounds, each O(nnz) scatter); space \(O(N)\) for the label array.

Backward

Non-differentiable: the output is a discrete integer labelling (no gradient flows through component assignment).

returns:

• labels (torch.Tensor) – Component label for each node, shape [N] (or [*batch, N] for batched input). Labels are in range [0, num_components), long, on self.device.

• num_components (int) – Number of connected components. For batched input every batch item shares the same sparsity pattern, so the partition is identical across the batch and a single int is returned.

Notes

• Matrix is treated as undirected (edges in either direction count)

• Self-loops are ignored for connectivity

• Batched: all batch items share row/col indices (same structure), so the component partition is the same for every batch item. The returned labels is broadcast to shape [*batch, N].

Examples

>>> import torch
>>> from torch_sla import SparseTensor
>>> # Edge 0-1 and isolated node 2  ->  components {0,1} and {2}
>>> row = torch.tensor([0, 1])
>>> col = torch.tensor([1, 0])
>>> val = torch.ones(2)
>>> A = SparseTensor(val, row, col, (3, 3))
>>> labels, num_comp = A.connected_components()
>>> num_comp
2
>>> labels.tolist()
[0, 0, 1]

has_isolated_components(*args, **kwargs)

to_connected_components(*args, **kwargs)

__matmul__(*args, **kwargs)

Matrix product A @ other (matvec / SpMM / sparse-sparse).

\[y = A x \quad(\text{vector}),\qquad Y = A X \quad(\text{dense block}),\qquad C = A B \quad(\text{sparse-sparse})\]

Algorithm

For a dense right operand, an index-scatter SpMV/SpMM accumulates each stored entry’s contribution; sparse-sparse routes to a gather-scatter sparse-times-sparse product. Vectorized, no Python per-nonzero loop:

y = zeros(M)
# one fused scatter-add over all nonzeros:
y[row[k]] += val[k] * x[col[k]]   for k in 0..nnz-1

Complexity

Time \(O(nnz)\) for a vector RHS (O(nnz \cdot c) for a dense block of c columns); space \(O(n)\) for the output.

Backward

Differentiable: \(\partial L/\partial x = A^{\top}(\partial L/\partial y)\) (another SpMV) and \(\partial L/\partial A_{rc} = (\partial L/\partial y)_r\, x_c\); \(O(1)\) extra graph nodes.

param other:

Right-hand operand: a dense vector [N] / matrix [N, c], or another SparseTensor for a sparse-sparse product.

type other:

torch.Tensor or SparseTensor

returns:

A @ other: dense for a dense operand, sparse for a sparse one.

rtype:

torch.Tensor or SparseTensor

Examples

>>> import torch
>>> from torch_sla import SparseTensor
>>> # A = diag(2, 3, 4);  A @ [1,1,1] = [2,3,4]
>>> d = torch.tensor([2.0, 3.0, 4.0])
>>> idx = torch.arange(3)
>>> A = SparseTensor(d, idx, idx, (3, 3))
>>> A @ torch.ones(3)
tensor([2., 3., 4.])

solve(*args, **kwargs)

Solve the sparse linear system \(Ax = b\).

\[A x = b \quad\Longrightarrow\quad x = A^{-1} b\]

Automatically handles batched tensors: if A is […batch, M, N] and b is […batch, M], returns x with shape […batch, N].

Algorithm

Dispatches (via torch_sla.linear_solve.spsolve()) to either a direct factorization (sparse LU / Cholesky / LDLᵀ) or a Krylov iterative method (CG for SPD, MINRES for symmetric-indefinite, BiCGStab / GMRES for general). method="auto" inspects symmetry and positive-definiteness to pick one. Krylov CG sketch:

r = b - A @ x;  p = r;  rs = <r, r>
repeat until ||r|| <= atol:
    Ap   = A @ p                 # one sparse matvec / iteration
    alpha = rs / <p, Ap>
    x    += alpha * p
    r    -= alpha * Ap
    rs_new = <r, r>
    p    = r + (rs_new / rs) * p
    rs   = rs_new

Complexity

Krylov: time \(O(m \cdot nnz)\) for m iterations, space \(O(n + nnz)\). Direct factorization: time between \(O(n^{1.5})\) (2-D) and \(O(n^{2})\) (3-D), space \(O(n\log n)\) to \(O(n^{4/3})\) for the fill-in factors.

Backward

Differentiable via the adjoint method: the backward pass solves \(A^{H}\lambda = \partial L/\partial x\) (reusing the same factorization / Krylov operator) and accumulates \(\partial L/\partial A = -\lambda\, x^{H}\). This adds only \(O(1)\) autograd graph nodes regardless of the iteration count m (one extra solve, not m recorded matvecs).

param b:

Right-hand side vector(s). Shape: - Non-batched: [M] or [M, K] for multiple RHS - Batched: […batch, M] or […batch, M, K]

type b:

torch.Tensor

param backend:

Solver backend. Default: “auto” (selects based on device). - “scipy”: Uses SciPy’s sparse solvers (CPU only) - “pytorch”: PyTorch-native iterative solvers (CPU & CUDA) - “cudss”: Uses NVIDIA cuDSS (CUDA only)

type backend:

{“auto”, “scipy”, “pytorch”, “cudss”}, optional

param method:

Solver method. Default: “auto” (selects based on matrix properties). - Direct methods: “lu”, “umfpack”, “cholesky”, “ldlt” - Iterative methods: “cg”, “bicgstab”, “gmres”, “minres”

type method:

str, optional

param atol:

Absolute tolerance for iterative solvers. Default: 1e-10.

type atol:

float, optional

param maxiter:

Maximum iterations for iterative solvers. Default: 10000.

type maxiter:

int, optional

param tol:

Relative tolerance for direct solvers. Default: 1e-12.

type tol:

float, optional

param verbose:

If True, print a one-line summary of the auto-selected backend, method, and detected matrix properties (symmetric, SPD). Default: False.

type verbose:

bool, optional

returns:

Solution x with same batch shape as b.

rtype:

torch.Tensor

raises ValueError:

If matrix is not square.

raises NotImplementedError:

If block sparse tensors are used (not yet supported).

Examples

>>> import torch
>>> from torch_sla import SparseTensor
>>> # 3x3 SPD matrix  A = diag(2,3,4)  (stored as 3 diagonal entries)
>>> val = torch.tensor([2.0, 3.0, 4.0])
>>> row = torch.tensor([0, 1, 2])
>>> col = torch.tensor([0, 1, 2])
>>> A = SparseTensor(val, row, col, (3, 3))
>>> b = torch.tensor([2.0, 3.0, 4.0])
>>> A.solve(b)
tensor([1., 1., 1.])

>>> # Batched solve: A is [4, 3, 3], b is [4, 3] -> x is [4, 3]
>>> A_batch = SparseTensor(val.expand(4, 3).clone(), row, col, (4, 3, 3))
>>> x_batch = A_batch.solve(b.expand(4, 3).clone())
>>> x_batch.shape
torch.Size([4, 3])

>>> # Specify backend / method explicitly
>>> x = A.solve(b, backend='scipy', method='cg')

solve_batch(*args, **kwargs)

Solve with different values but same sparsity structure.

\[A_i x_i = b_i,\quad i = 1\ldots B \qquad\text{with } \mathrm{pattern}(A_i)\equiv(\text{row},\text{col})\]

This is efficient when you have the same structure but different values (e.g., time-stepping, optimization, parameter sweeps): the symbolic factorization / sparsity analysis is shared across the batch.

Algorithm

Detect symmetry/SPD once from the first batch element, then loop over the B value-vectors, calling spsolve() per system (each reuses the same row/col indices):

props = analyze(values[0], row, col)      # symmetry / SPD, once
for i in 1..B:
    x[i] = spsolve(values[i], row, col, b[i], props)

Complexity

B independent solves: time \(O(B \cdot m \cdot nnz)\) (Krylov) or \(O(B \cdot n^{1.5..2})\) (direct); space \(O(n + nnz)\) per system. Backward mirrors solve() (one adjoint solve per system, \(O(1)\) extra graph nodes each).

param values:

Matrix values. Shape […batch, nnz] where … are batch dimensions. All matrices share the same row_indices and col_indices.

type values:

torch.Tensor

param b:

Right-hand side. Shape […batch, M].

type b:

torch.Tensor

param backend:

Solver backend. See solve() for details. Default: “auto”.

type backend:

{“auto”, “scipy”, “pytorch”, “cudss”}, optional

param method:

Solver method. See solve() for details. Default: “auto”.

type method:

str, optional

param atol:

Absolute tolerance for iterative solvers. Default: 1e-10.

type atol:

float, optional

param maxiter:

Maximum iterations for iterative solvers. Default: 10000.

type maxiter:

int, optional

param tol:

Relative tolerance. Default: 1e-12.

type tol:

float, optional

returns:

Solution x with shape […batch, N].

rtype:

torch.Tensor

Examples

>>> # Template matrix
>>> A = SparseTensor(val, row, col, (10, 10))

>>> # Batch of different values
>>> val_batch = torch.stack([val * (1 + 0.1*i) for i in range(4)])  # [4, nnz]
>>> b_batch = torch.randn(4, 10)

>>> # Solve all at once
>>> x_batch = A.solve_batch(val_batch, b_batch)  # [4, 10]

nonlinear_solve(*args, **kwargs)

Solve nonlinear equation \(F(u, A, \theta) = 0\) with adjoint gradients.

\[F(u^\*, A, \theta) = 0,\qquad u_{n+1} = u_n - J(u_n)^{-1} F(u_n),\quad J = \frac{\partial F}{\partial u}\]

The SparseTensor A is automatically passed as the first parameter to the residual function, enabling gradients to flow through A’s values.

Algorithm

Newton-Raphson: at each step assemble the sparse Jacobian \(J = \partial F/\partial u\) (explicit jac_fn or torch.autograd.functional.jacobian) and solve the linear update via spsolve(); an optional Armijo line search damps the step.

u = u0
for it in 1..max_iter:
    F = residual(u, A, *params)
    if ||F|| <= atol or ||F|| <= tol*||F0||: break
    J  = jacobian(F, u)            # sparse dF/du
    du = spsolve(J, -F)            # Newton direction
    a  = armijo_line_search(u, du) # if line_search
    u  = u + a * du

Complexity

Time \(O(m \cdot \text{solve})\) for m Newton steps, where each solve is one sparse linear solve (Krylov \(O(m'\,nnz)\) or direct); space \(O(\text{solve})\), i.e. that of a single linear solve plus the Jacobian.

Backward

Gradients use the implicit-function theorem (not unrolled Newton): backward solves one adjoint system \(J^{H}\lambda = \partial L/\partial u\) at the converged \(u^\*\) and back-propagates \(\partial L/\partial(\cdot) = -\lambda^{H}\,\partial F/\partial(\cdot)\). This is \(O(1)\) extra graph nodes regardless of m.

param residual_fn:

Function F(u, A, *params) -> residual tensor. - u: Current solution estimate - A: This SparseTensor (passed automatically) - *params: Additional parameters with requires_grad=True

type residual_fn:

Callable

param u0:

Initial guess for solution.

type u0:

torch.Tensor

param *params:

Additional parameters (e.g., boundary conditions, coefficients). Tensors with requires_grad=True will receive gradients.

type *params:

torch.Tensor

param jac_fn:

Optional explicit Jacobian J(u, A, *params) -> (val, row, col, shape) returning the sparse dF/du in COO form. If None (default) the Jacobian is obtained via torch.autograd.functional.jacobian.

type jac_fn:

Callable, optional

param method:

Nonlinear solver method. Only Newton-Raphson (with implicit-diff gradients) is supported; each step solves the sparse Jacobian system via spsolve.

type method:

{‘newton’}, optional

param tol:

Relative convergence tolerance on ||F||. Default: 1e-8.

type tol:

float, optional

param atol:

Absolute convergence tolerance on ||F||. Default: 1e-12.

type atol:

float, optional

param max_iter:

Maximum nonlinear iterations. Default: 50.

type max_iter:

int, optional

param line_search:

Use Armijo line search for Newton. Default: True.

type line_search:

bool, optional

param verbose:

Print convergence information. Default: False.

type verbose:

bool, optional

param linear_solver:

Backend for the linear (Jacobian) solves. Default: ‘auto’.

type linear_solver:

str, optional

param linear_method:

Method for the linear (Jacobian) solves. Default: ‘auto’. Note: the Jacobian of a general nonlinear residual is NOT symmetric, so a direct method (e.g. ‘lu’) is recommended over ‘cg’.

type linear_method:

str, optional

returns:

Solution u* satisfying F(u*, A, θ) ≈ 0.

rtype:

torch.Tensor

Examples

>>> # Nonlinear PDE: A @ u + u² = f
>>> def residual(u, A, f):
...     return A @ u + u**2 - f
...
>>> A = SparseTensor(val, row, col, (n, n))
>>> f = torch.randn(n, requires_grad=True)
>>> u0 = torch.zeros(n)
>>>
>>> u = A.nonlinear_solve(residual, u0, f, method='newton')
>>>
>>> # Gradients flow via adjoint method
>>> loss = u.sum()
>>> loss.backward()
>>> print(f.grad)  # ∂u/∂f
>>> print(A.values.grad)  # ∂u/∂A (if A.values.requires_grad)

>>> # Nonlinear elasticity: K(u) @ u = F
>>> def residual_elasticity(u, K, F, material):
...     # K depends on displacement through material nonlinearity
...     return K @ u - F + material * u**3
...
>>> u = K.nonlinear_solve(residual_elasticity, u0, F, material)

norm(ord: Literal['fro', 1, 2] = 'fro') → Tensor

Compute a matrix norm.

\[\lVert A\rVert_F = \Big(\sum_{ij} |A_{ij}|^2\Big)^{1/2},\quad \lVert A\rVert_1 = \max_j \sum_i |A_{ij}|,\quad \lVert A\rVert_2 = \sigma_{\max}(A)\]

For batched tensors, returns a norm per batch element.

Algorithm

The Frobenius norm is just the 2-norm of the stored value vector (zeros contribute nothing). The 1-norm reduces over column sums of |val|; the 2-norm is the largest singular value (Lanczos).

'fro' -> sqrt(sum(val**2))          # over nonzeros only
1     -> max over columns of sum|val|
2     -> svd(A).singular_values.max()

Complexity

Frobenius / 1-norm: time \(O(nnz)\), space \(O(1)\). The 2-norm costs one Lanczos SVD, \(O(m\,nnz)\).

param ord:

Norm type: - ‘fro’: Frobenius norm (default) - 1: Maximum absolute column sum - 2: Spectral norm (largest singular value)

type ord:

{‘fro’, 1, 2}, optional

returns:

Norm value(s). Shape [] for non-batched, [*batch_shape] for batched.

rtype:

torch.Tensor

Examples

>>> import torch
>>> from torch_sla import SparseTensor
>>> # entries 3 and 4  ->  Frobenius norm = sqrt(9 + 16) = 5
>>> val = torch.tensor([3.0, 4.0])
>>> row = torch.tensor([0, 1])
>>> col = torch.tensor([0, 1])
>>> A = SparseTensor(val, row, col, (2, 2))
>>> A.norm('fro')
tensor(5.)

spy(*args, **kwargs)

Render the sparsity pattern as an image. See viz.spy().

\[\begin{split}\mathrm{image}[i, j] = \begin{cases} |A_{ij}| / \max|A| & A_{ij}\neq 0 \\ \text{white (NaN)} & A_{ij}=0 \end{cases}\end{split}\]

Algorithm

Scatter the stored entries into an M x N raster coloured by normalised magnitude, leaving structural zeros transparent, then imshow it:

image = full((M, N), NaN)
image[row, col] = |val| / max|val|
imshow(image)            # zeros render white

Complexity

Time \(O(nnz + MN)\) (scatter + raster allocation), space \(O(MN)\) for the image. Non-differentiable (plotting only).

returns:

The axes the pattern was drawn on.

rtype:

matplotlib.axes.Axes

Examples

>>> import torch
>>> from torch_sla import SparseTensor
>>> A = SparseTensor.tridiagonal(10)
>>> ax = A.spy()

eigs(*args, **kwargs)

Compute k eigenvalues and eigenvectors of a (general) sparse matrix.

\[A v_i = \lambda_i v_i,\qquad i = 1\ldots k\]

For batched tensors, computes for each batch element. Symmetric matrices (and which in {"LA","SA"}) are routed to the more efficient eigsh().

Algorithm

A subspace iteration (LOBPCG for symmetric / Hermitian operators; see _lobpcg_core()) projects A onto a small search space and repeatedly solves the tiny dense Rayleigh-Ritz problem:

X = orthonormal random block, columns ~ k
repeat until Ritz residual small:
    AX = A @ X                       # k sparse matvecs
    H  = X^H A X                     # small (m x m) Gram matrix
    (theta, C) = eigh(H)             # dense Rayleigh-Ritz
    X  = orthonormalize([X | (AX - X theta) | P])  # X | residual | conj-dir
return theta[:k], X[:k]

Complexity

Time \(O(m\,k\,nnz)\) (m outer iterations, k matvecs each); space \(O(k n + nnz)\) for the block of k Ritz vectors plus the matrix.

Backward

Differentiable via the adjoint method with \(\partial L/\partial A = \sum_i (\partial L/\partial\lambda_i)\, v_i v_i^{H}\) (eigenvalue part); \(O(1)\) extra graph nodes, independent of the iteration count. For complex eigenvalues only the real part is differentiable.

param k:

Number of eigenvalues to compute. Default: 6.

type k:

int, optional

param which:

Which eigenvalues to find: - “LM”: Largest magnitude (default) - “SM”: Smallest magnitude - “LR”/”SR”: Largest/smallest real part - “LA”/”SA”: Largest/smallest algebraic (for symmetric)

type which:

{“LM”, “SM”, “LR”, “SR”, “LA”, “SA”}, optional

param sigma:

Find eigenvalues near sigma (shift-invert mode).

type sigma:

float, optional

param return_eigenvectors:

Whether to return eigenvectors. Default: True.

type return_eigenvectors:

bool, optional

returns:

• eigenvalues (torch.Tensor) – Shape [k] for non-batched, [*batch_shape, k] for batched.

• eigenvectors (torch.Tensor or None) – Shape [M, k] for non-batched, [*batch_shape, M, k] for batched. None if return_eigenvectors is False.

Notes

Gradient Support:

• Both CPU and CUDA: Fully differentiable via adjoint method

• Uses O(1) graph nodes regardless of iteration count

• For symmetric matrices, prefer eigsh() for efficiency

Warning: For non-symmetric matrices with complex eigenvalues, gradient computation is only supported for the real part.

Examples

>>> import torch
>>> from torch_sla import SparseTensor
>>> # A = diag(1, 2, 3, 4): eigenvalues are the diagonal entries
>>> d = torch.tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
>>> idx = torch.arange(4)
>>> A = SparseTensor(d, idx, idx, (4, 4))
>>> evals, evecs = A.eigs(k=2, which="LM")  # two largest-magnitude
>>> evals.detach().sort().values
tensor([3., 4.])
>>> evals.real.sum().backward()  # differentiable w.r.t. A.values

eigsh(*args, **kwargs)

Compute k eigenvalues for symmetric / Hermitian matrices.

\[A v_i = \lambda_i v_i,\quad A = A^{H},\qquad i = 1\ldots k\]

More efficient than eigs() for symmetric matrices: it exploits \(A=A^{H}\) so all Ritz values are real and a single eigh per iteration suffices.

Algorithm

LOBPCG (Knyazev 2001; see _lobpcg_core()) – locally optimal block preconditioned conjugate gradient. Maintains a 3-block subspace [X | residual | conjugate-direction] and minimizes the Rayleigh quotient over it each step:

X = orthonormal random block (m >= k columns)
repeat until ||A x_i - lambda_i x_i|| <= tol*|lambda_i|:
    theta, X = rayleigh_ritz(A, [X | R | P])  # dense eigh on Gram
    R = A @ X - X * theta                     # block residual
    P = conjugate direction (Knyazev eq. 7)
return theta[:k], X[:k]

Complexity

Time \(O(m\,k\,nnz)\); space \(O(k n + nnz)\).

Backward

Adjoint method: \(\partial L/\partial A = \sum_i (\partial L/\partial\lambda_i)\, v_i v_i^{\top}\); \(O(1)\) extra graph nodes regardless of the iteration count.

param k:

Number of eigenvalues to compute. Default: 6.

type k:

int, optional

param which:

Which eigenvalues to find: - “LM”: Largest magnitude (default) - “SM”: Smallest magnitude - “LA”/”SA”: Largest/smallest algebraic

type which:

{“LM”, “SM”, “LA”, “SA”}, optional

param sigma:

Find eigenvalues near sigma.

type sigma:

float, optional

param return_eigenvectors:

Whether to return eigenvectors. Default: True.

type return_eigenvectors:

bool, optional

returns:

• eigenvalues (torch.Tensor) – Shape [k] for non-batched, [*batch_shape, k] for batched.

• eigenvectors (torch.Tensor or None) – Shape [M, k] for non-batched, [*batch_shape, M, k] for batched.

Notes

Device support: CPU and CUDA are first-class. MPS is not recommended: PyTorch’s MPS backend forces float32 (caps Ritz residual at ~1e-4..1e-3 on PDE-like operators) and is missing native linalg.eigh / efficient tall-skinny linalg.qr kernels, so the core internally round-trips both to CPU as a workaround. The fallback works but most of the LOBPCG work actually happens on CPU; a RuntimeWarning fires on every MPS call to point users at device='cpu' or 'cuda'.

Gradient Support:

• Both CPU and CUDA: Fully differentiable via adjoint method

• Uses O(1) graph nodes regardless of iteration count

• Gradient computed as: ∂L/∂A = Σ_i (∂L/∂λ_i) * v_i @ v_i.T

Examples

>>> import torch
>>> from torch_sla import SparseTensor
>>> # Symmetric A = diag(1, 2, 3, 4)
>>> d = torch.tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
>>> idx = torch.arange(4)
>>> A = SparseTensor(d, idx, idx, (4, 4))
>>> evals, evecs = A.eigsh(k=2, which="SA")  # two smallest algebraic
>>> evals.detach().sort().values
tensor([1., 2.])
>>> evals.sum().backward()   # computes d(loss)/d(A.values)

svd(*args, **kwargs)

Compute truncated (rank-k) singular value decomposition.

\[A \approx U \Sigma V^{\top},\qquad \Sigma = \mathrm{diag}(\sigma_1 \ge \cdots \ge \sigma_k \ge 0)\]

Algorithm

A Lanczos bidiagonalization (scipy.sparse.linalg.svds / ARPACK on CPU) builds a Krylov subspace from sparse matvecs against A and Aᵀ and extracts the k dominant singular triplets from the small bidiagonal factor:

build Krylov basis via alternating  v -> A v,  u -> A^T u
B = bidiagonal projection of A onto that basis
(U_b, S, V_b) = dense_svd(B)         # small problem
U = Q_left @ U_b;  V = Q_right @ V_b
return U[:, :k], S[:k], V[:, :k]^T

Complexity

Time \(O(m\,k\,nnz)\) (m Lanczos steps); space \(O(k n + nnz)\).

Backward

Adjoint method: the singular-value gradient flows through the stored pattern as \(\partial L/\partial A_{rc} = \sum_i (\partial L/\partial\sigma_i)\, U_{r i} V_{c i}\); \(O(1)\) extra graph nodes.

param k:

Number of singular values to compute. Default: 6.

type k:

int, optional

returns:

• U (torch.Tensor) – Left singular vectors. Shape [M, k] or [*batch_shape, M, k].

• S (torch.Tensor) – Singular values. Shape [k] or [*batch_shape, k].

• Vt (torch.Tensor) – Right singular vectors. Shape [k, N] or [*batch_shape, k, N].

Notes

Gradient Support:

• CUDA: not currently supported (raises NotImplementedError; move to CPU first).

• CPU: singular values are differentiable via the adjoint above; the singular vectors break the gradient chain (SciPy forward).

For fully differentiable SVD on CPU, use A.to_dense() and torch.linalg.svd().

Examples

>>> import torch
>>> from torch_sla import SparseTensor
>>> # A = diag(4, 1, 3): singular values are |diagonal|, sorted desc
>>> d = torch.tensor([4.0, 1.0, 3.0])
>>> idx = torch.arange(3)
>>> A = SparseTensor(d, idx, idx, (3, 3))
>>> U, S, Vt = A.svd(k=2)
>>> S.detach().round().sort().values
tensor([3., 4.])

condition_number(*args, **kwargs)

Estimate the condition number \(\kappa(A)\).

\[\kappa_2(A) = \frac{\sigma_{\max}(A)}{\sigma_{\min}(A)},\qquad \kappa_p(A) = \lVert A\rVert_p\,\lVert A^{-1}\rVert_p\]

Algorithm

For ord=2 take the ratio of the extreme singular values from a truncated svd() (Lanczos). For other orders, estimate \(\lVert A^{-1}\rVert\) from a single linear solve against a random unit vector:

if ord == 2:
    S = svd(A).singular_values
    return S.max() / S.min()
else:
    e = random_unit_vector()
    x = A.solve(e)                  # ~ ||A^{-1}|| estimate
    return ||A||_ord * ||x|| / ||e||

Complexity

Time \(O(m\,nnz)\) (one Lanczos SVD or one Krylov solve), space \(O(n + nnz)\).

param ord:

Norm order for condition number. Default: 2 (spectral).

type ord:

int, optional

returns:

Condition number. Shape [] or [*batch_shape].

rtype:

torch.Tensor

Examples

>>> import torch
>>> from torch_sla import SparseTensor
>>> # A = diag(1, 2, 4): kappa_2 = sigma_max / sigma_min = 4 / 1
>>> d = torch.tensor([1.0, 2.0, 4.0])
>>> idx = torch.arange(3)
>>> A = SparseTensor(d, idx, idx, (3, 3))
>>> A.condition_number(ord=2).round()
tensor(4.)

det(*args, **kwargs)

Compute determinant of the sparse matrix with gradient support.

\[\det(A) = \prod_i U_{ii}\quad\text{from } PA = LU, \qquad U_{ii} = \text{pivots}\]

Uses LU decomposition (CPU) or dense conversion (CUDA) to compute the determinant efficiently. Supports automatic differentiation via the adjoint method.

Algorithm

Sparse LU factorize \(PA = LU\) and multiply the pivots (L is unit-diagonal), correcting the sign by the permutation parity:

P, L, U = sparse_lu(A)
det = sign(P) * prod(diag(U))

Complexity

Time \(O(n^{1.5})\) (2-D fill) to \(O(n^{2})\) (3-D); space \(O(n\log n)\) to \(O(n^{4/3})\) for the LU factors.

Backward

Adjoint via Jacobi’s formula: \(\partial\det(A)/\partial A = \det(A)\,(A^{-\top})\), i.e. the backward reuses the factorization for one transposed solve; \(O(1)\) extra graph nodes.

returns:

Determinant value. Shape [] for single matrix or [*batch_shape] for batched.

rtype:

torch.Tensor

raises ValueError:

If matrix is not square

Notes

• Only square matrices have determinants

• For large matrices, determinant values can overflow/underflow

• Consider using log-determinant for numerical stability in such cases

• Supports both CPU (via SciPy) and CUDA (via torch.linalg.det)

• For batched tensors, computes determinant independently for each batch

• Fully differentiable: gradients computed via adjoint method

• Gradient formula: ∂det(A)/∂A = det(A) * (A^{-1})^T

Performance Warning

CUDA performance is significantly slower than CPU for sparse matrices!

• CPU: Uses sparse LU decomposition (O(nnz^1.5)), ~0.3-0.8ms for n=10-1000

• CUDA: Converts to dense (O(n²) memory + O(n³) compute), ~0.2-2.5ms

The CUDA version requires converting the sparse matrix to dense format because cuDSS doesn’t expose determinant computation for sparse matrices. This makes it inefficient for large sparse matrices.

Recommendation: For sparse matrices, use .cpu().det().cuda() instead:

>>> # Slow: CUDA with dense conversion
>>> det_slow = A_cuda.det()  # ~2.5ms for n=1000
>>>
>>> # Fast: CPU with sparse LU
>>> det_fast = A_cuda.cpu().det()  # ~0.8ms for n=1000
>>> det_fast = det_fast.cuda()  # Move result back if needed

Examples

>>> # Simple 2x2 matrix
>>> val = torch.tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
>>> row = torch.tensor([0, 0, 1, 1])
>>> col = torch.tensor([0, 1, 0, 1])
>>> A = SparseTensor(val, row, col, (2, 2))
>>> det = A.det()
>>> print(det)  # Should be -2.0
>>> det.backward()
>>> print(val.grad)  # Gradient w.r.t. matrix values
>>>
>>> # CUDA support
>>> A_cuda = A.cuda()
>>> det_cuda = A_cuda.det()
>>>
>>> # Batched matrices
>>> val_batch = val.unsqueeze(0).expand(3, -1).clone()
>>> A_batch = SparseTensor(val_batch, row, col, (3, 2, 2))
>>> det_batch = A_batch.det()
>>> print(det_batch.shape)  # torch.Size([3])

logdet(*args, **kwargs)

Log-determinant \(\log|\det A|\) of this matrix. See torch_sla.det.

\[\log\lvert\det A\rvert = \sum_i \log\lvert U_{ii}\rvert \;=\; \mathrm{tr}\,\log A \;\approx\; \frac{1}{p}\sum_{j=1}^{p} z_j^{\top}\log(A)\,z_j\]

Numerically stable alternative to det() (no overflow): works in log-space directly.

Algorithm

Two regimes. Direct (small / general): factorize and sum log|diag(U)| (or 2*sum log diag(L) for Cholesky on SPD). Stochastic (large SPD, default method='auto'): the Hutchinson estimator approximates \(\mathrm{tr}\log A\) with matvec-only probes (no factorization, distributed-friendly):

# direct
P, L, U = sparse_lu(A);  return sum(log|diag(U)|)
# stochastic (Hutchinson + matvec log(A) z)
for j in 1..p:  acc += z_j^T (log A) z_j   # z_j Rademacher probes
return acc / p

Complexity

Direct: time \(O(n^{1.5})\)–\(O(n^{2})\), space \(O(n\log n)\)–\(O(n^{4/3})\). Stochastic: dominated by p matvec sequences, \(O(p\,m\,nnz)\) time, \(O(n+nnz)\) space.

Backward

Adjoint \(\partial \log\lvert\det A\rvert/\partial A = A^{-\top}\); \(O(1)\) extra graph nodes.

Examples

>>> import torch
>>> from torch_sla import SparseTensor
>>> # A = diag(1, e, e^2): logdet = 0 + 1 + 2 = 3
>>> import math
>>> d = torch.tensor([1.0, math.e, math.e ** 2])
>>> idx = torch.arange(3)
>>> A = SparseTensor(d, idx, idx, (3, 3))
>>> A.logdet().round()
tensor(3.)

lu(*args, **kwargs)

Compute the LU decomposition for repeated solves.

\[P A Q = L U\]

with L unit-lower-triangular, U upper-triangular, and P/Q row/column permutations chosen to limit fill-in.

Algorithm

Sparse Gaussian elimination with a fill-reducing reordering (SuperLU / UMFPACK via SciPy). The returned object caches the factors so each subsequent solve() is just two cheap triangular sweeps:

Q = fill_reducing_order(A)          # symbolic
P, L, U = factorize(A[:, Q])        # numeric
# later, per RHS:
solve(b): y = L \\ (P b);  x = Q (U \\ y)

Complexity

Factorization (once): time \(O(n^{1.5})\)–\(O(n^{2})\), space \(O(n\log n)\)–\(O(n^{4/3})\) for the factors. Each reuse: \(O(nnz(L)+nnz(U))\) per right-hand side.

returns:

Factorization object with a solve() method.

rtype:

LUFactorization

Examples

>>> import torch
>>> from torch_sla import SparseTensor
>>> d = torch.tensor([2.0, 4.0])
>>> idx = torch.arange(2)
>>> A = SparseTensor(d, idx, idx, (2, 2))
>>> lu = A.lu()
>>> lu.solve(torch.tensor([2.0, 4.0]))
tensor([1., 1.])
>>> lu.solve(torch.tensor([4.0, 8.0]))   # reuses the factorization
tensor([2., 2.])

sum(*args, **kwargs)

mean(*args, **kwargs)

prod(*args, **kwargs)

max(*args, **kwargs)

min(*args, **kwargs)

abs(*args, **kwargs)

sqrt(*args, **kwargs)

square(*args, **kwargs)

exp(*args, **kwargs)

log(*args, **kwargs)

save(path: str | PathLike, metadata: Dict[str, str] | None = None) → None

Save SparseTensor to safetensors format.

Parameters:

• path (str or PathLike) – Output file path (should end with .safetensors).

• metadata (dict, optional) – Additional metadata to store.

Example

>>> A = SparseTensor(val, row, col, (100, 100))
>>> A.save("matrix.safetensors")

classmethod load(path: str | PathLike, device: str | device = 'cpu') → SparseTensor

Load SparseTensor from safetensors format.

Parameters:

• path (str or PathLike) – Input file path.

• device (str or torch.device) – Device to load tensors to.

Returns:

The loaded sparse tensor.

Return type:

SparseTensor

Example

>>> A = SparseTensor.load("matrix.safetensors", device="cuda")

SparseTensorList

Container for multiple sparse matrices with different sparsity patterns. Useful for batched operations on heterogeneous graphs.

class torch_sla.SparseTensorList(tensors: List[SparseTensor])

Bases: object

A list of SparseTensors with different structures.

Provides a unified interface for batch operations on matrices with different sparsity patterns. Unlike a batched SparseTensor (which requires every matrix to share one structure), SparseTensorList lets each matrix have its own shape and sparsity pattern, then dispatches solve / property checks / elementwise ops over the whole collection.

You typically get one of these in two ways:

• Construct it directly from a Python list of SparseTensor (or via from_coo_list() / from_torch_sparse_list()) when you have several independently-patterned matrices to solve together.

• As the return value of SparseTensor.to_connected_components(), which splits a block-diagonal matrix into one SparseTensor per connected component.

Parameters:

tensors (List[SparseTensor]) – List of SparseTensor objects.

shapes

List of shapes for each tensor.

Type:

List[Tuple[int, …]]

device

Device (from first tensor).

Type:

torch.device

dtype

Data type (from first tensor).

Type:

torch.dtype

Examples

>>> # Create matrices with different sizes
>>> A1 = SparseTensor(val1, row1, col1, (10, 10))
>>> A2 = SparseTensor(val2, row2, col2, (20, 20))
>>> A3 = SparseTensor(val3, row3, col3, (30, 30))

>>> # Create list
>>> matrices = SparseTensorList([A1, A2, A3])
>>> print(matrices.shapes)  # [(10, 10), (20, 20), (30, 30)]

>>> # Batch solve
>>> x_list = matrices.solve([b1, b2, b3])

>>> # Check properties for all
>>> is_sym = matrices.is_symmetric()  # [tensor(True), tensor(True), tensor(True)]

classmethod from_coo_list(matrices: List[Tuple[Tensor, Tensor, Tensor, Tuple[int, ...]]]) → SparseTensorList

Create from list of COO data tuples.

Parameters:

matrices (List[Tuple]) – List of (values, row_indices, col_indices, shape) tuples.

Returns:

List of SparseTensors.

Return type:

SparseTensorList

Examples

>>> data = [
...     (val1, row1, col1, (10, 10)),
...     (val2, row2, col2, (20, 20)),
... ]
>>> matrices = SparseTensorList.from_coo_list(data)

classmethod from_torch_sparse_list(A_list: List[Tensor]) → SparseTensorList

Create from list of PyTorch sparse tensors.

Parameters:

A_list (List[torch.Tensor]) – List of PyTorch sparse COO tensors.

Returns:

List of SparseTensors.

Return type:

SparseTensorList

property shapes: List[Tuple[int, ...]]

List of shapes for each tensor.

property device: device

Device of the first tensor.

property dtype: dtype

Data type of the first tensor.

to(device: str | device) → SparseTensorList

Move all tensors to device.

Parameters:

device (str or torch.device) – Target device.

Returns:

New list with tensors on target device.

Return type:

SparseTensorList

cuda() → SparseTensorList

Move all tensors to CUDA.

cpu() → SparseTensorList

Move all tensors to CPU.

sum(axis: int | None = None) → List[Tensor] | Tensor

Sum values in each matrix.

Parameters:

axis (int, optional) – If None: sum all values in each matrix, return List[scalar]. If 0: sum over rows for each matrix. If 1: sum over columns for each matrix.

Returns:

If axis is None: List of scalar tensors (one per matrix). If axis is 0 or 1: List of 1D tensors.

Return type:

List[torch.Tensor] or torch.Tensor

Examples

>>> matrices = SparseTensorList([A1, A2, A3])
>>> totals = matrices.sum()  # [sum(A1), sum(A2), sum(A3)]
>>> row_sums = matrices.sum(axis=1)  # [A1.sum(1), A2.sum(1), ...]

mean(axis: int | None = None) → List[Tensor]

Mean of values in each matrix.

Parameters:

axis (int, optional) – Same as sum().

Returns:

List of mean values/vectors.

Return type:

List[torch.Tensor]

max() → List[Tensor]

Maximum value in each matrix.

min() → List[Tensor]

Minimum value in each matrix.

abs() → SparseTensorList

Absolute value of all elements.

clamp(min: float | None = None, max: float | None = None) → SparseTensorList

Clamp values in all matrices.

pow(exponent: float) → SparseTensorList

Element-wise power.

sqrt() → SparseTensorList

Element-wise square root.

exp() → SparseTensorList

Element-wise exponential.

log() → SparseTensorList

Element-wise natural logarithm.

solve(b_list: List[Tensor], **kwargs) → List[Tensor]

Solve linear systems for all matrices.

Parameters:

• b_list (List[torch.Tensor]) – List of right-hand side vectors, one per matrix.

• **kwargs – Additional arguments passed to SparseTensor.solve().

Returns:

List of solutions.

Return type:

List[torch.Tensor]

Examples

>>> matrices = SparseTensorList([A1, A2, A3])
>>> x_list = matrices.solve([b1, b2, b3])

is_symmetric(**kwargs) → List[Tensor]

Check symmetry for all matrices.

Parameters:

**kwargs – Arguments passed to SparseTensor.is_symmetric().

Returns:

List of boolean tensors.

Return type:

List[torch.Tensor]

is_positive_definite(**kwargs) → List[Tensor]

Check positive definiteness for all matrices.

Parameters:

**kwargs – Arguments passed to SparseTensor.is_positive_definite().

Returns:

List of boolean tensors.

Return type:

List[torch.Tensor]

norm(ord: Literal['fro', 1, 2] = 'fro') → List[Tensor]

Compute norms for all matrices.

Parameters:

ord ({'fro', 1, 2}) – Norm type.

Returns:

List of norm values.

Return type:

List[torch.Tensor]

eigs(k: int = 6, **kwargs) → List[Tuple[Tensor, Tensor | None]]

Compute eigenvalues for all matrices.

Parameters:

• k (int) – Number of eigenvalues.

• **kwargs – Additional arguments.

Returns:

List of (eigenvalues, eigenvectors) tuples.

Return type:

List[Tuple[torch.Tensor, Optional[torch.Tensor]]]

eigsh(k: int = 6, **kwargs) → List[Tuple[Tensor, Tensor | None]]

Compute eigenvalues for symmetric matrices.

Parameters:

• k (int) – Number of eigenvalues.

• **kwargs – Additional arguments.

Returns:

List of (eigenvalues, eigenvectors) tuples.

Return type:

List[Tuple[torch.Tensor, Optional[torch.Tensor]]]

svd(k: int = 6) → List[Tuple[Tensor, Tensor, Tensor]]

Compute SVD for all matrices.

Parameters:

k (int) – Number of singular values.

Returns:

List of (U, S, Vt) tuples.

Return type:

List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]

condition_number(ord: int = 2) → List[Tensor]

Compute condition numbers for all matrices.

Parameters:

ord (int) – Norm order.

Returns:

List of condition numbers.

Return type:

List[torch.Tensor]

det() → List[Tensor]

Compute determinants for all matrices.

Returns:

List of determinant values.

Return type:

List[torch.Tensor]

Examples

>>> matrices = SparseTensorList([A1, A2, A3])
>>> dets = matrices.det()
>>> print([d.item() for d in dets])

spy(indices: List[int] | None = None, ncols: int = 3, figsize: Tuple[float, float] | None = None, **kwargs)

Visualize sparsity patterns for multiple matrices in a grid.

Parameters:

• indices (List[int], optional) – Which matrices to visualize. Default: all.

• ncols (int, optional) – Number of columns in subplot grid. Default: 3.

• figsize (Tuple[float, float], optional) – Figure size. Auto-computed if None.

• **kwargs – Additional arguments passed to SparseTensor.spy().

Returns:

fig – The figure object.

Return type:

matplotlib.figure.Figure

Examples

>>> matrices = SparseTensorList([A1, A2, A3, A4])
>>> matrices.spy()  # Visualize all in grid
>>> matrices.spy(indices=[0, 2])  # Visualize specific ones

to_block_diagonal() → SparseTensor

Merge all matrices into a single block-diagonal SparseTensor.

Creates a sparse matrix where each input matrix appears as a block on the diagonal: diag(A1, A2, …, An).

Returns:

Block-diagonal matrix with shape (sum(M_i), sum(N_i)).

Return type:

SparseTensor

Notes

The resulting matrix has the structure:

` [A1  0  0  ...] [ 0 A2  0  ...] [ 0  0 A3  ...] [... ... ... ] `

Examples

>>> A1 = SparseTensor(val1, row1, col1, (10, 10))
>>> A2 = SparseTensor(val2, row2, col2, (20, 20))
>>> stl = SparseTensorList([A1, A2])
>>> A_block = stl.to_block_diagonal()  # Shape (30, 30)

classmethod from_block_diagonal(sparse: SparseTensor, sizes: List[Tuple[int, int]]) → SparseTensorList

Split a block-diagonal SparseTensor into a list of matrices.

Parameters:

• sparse (SparseTensor) – Block-diagonal matrix to split.

• sizes (List[Tuple[int, int]]) – List of (rows, cols) for each block. Must sum to sparse.shape.

Returns:

List of extracted blocks.

Return type:

SparseTensorList

Examples

>>> A_block = SparseTensor(val, row, col, (30, 30))
>>> stl = SparseTensorList.from_block_diagonal(A_block, [(10, 10), (20, 20)])
>>> print(len(stl))  # 2

property block_sizes: List[Tuple[int, int]]

Get the (rows, cols) size of each matrix.

Returns:

List of (M, N) tuples.

Return type:

List[Tuple[int, int]]

property total_nnz: int

Total number of non-zeros across all matrices.

property total_shape: Tuple[int, int]

Shape of the block-diagonal representation.

LUFactorization

LU factorization for efficient repeated solves with the same matrix.

class torch_sla.LUFactorization(lu_factor, shape: Tuple[int, int], dtype: dtype, device: device)

Bases: object

LU factorization wrapper for efficient repeated solves.

Created by SparseTensor.lu().

Parameters:

• lu_factor (scipy.sparse.linalg.SuperLU) – The SciPy LU factorization object.

• shape (Tuple[int, int]) – Matrix shape.

• dtype (torch.dtype) – Data type.

• device (torch.device) – Device.

Examples

>>> A = SparseTensor(val, row, col, (10, 10))
>>> lu = A.lu()
>>> x1 = lu.solve(b1)  # First solve
>>> x2 = lu.solve(b2)  # Much faster - reuses factorization

solve(b: Tensor) → Tensor

Solve Ax = b using the cached factorization.

Parameters:

b (torch.Tensor) – Right-hand side vector.

Returns:

Solution x.

Return type:

torch.Tensor

---

Distributed Classes

DSparseTensor

Distributed sparse tensor with domain decomposition support. Uses halo exchange for communication between partitions.

class torch_sla.DSparseTensor

Bases: object

Distributed Sparse Tensor with automatic partitioning and halo exchange.

A Pythonic wrapper that provides a unified interface for distributed sparse matrix operations. Supports indexing to access individual partitions.

Parameters:

• values (torch.Tensor) – Non-zero values [nnz]

• row_indices (torch.Tensor) – Row indices [nnz]

• col_indices (torch.Tensor) – Column indices [nnz]

• shape (Tuple[int, int]) – Matrix shape (m, n)

• num_partitions (int) – Number of partitions to create

• coords (torch.Tensor, optional) – Node coordinates for geometric partitioning [num_nodes, dim]

• partition_method (str) – Partitioning method: ‘metis’, ‘rcb’, ‘slicing’, ‘simple’

• device (str or torch.device) – Device for the matrix data

• verbose (bool) – Whether to print partition info

Example

>>> import torch
>>> from torch_sla import DSparseTensor
>>>
>>> # Create distributed tensor with 4 partitions
>>> A = DSparseTensor(val, row, col, shape, num_partitions=4)
>>>
>>> # Access individual partitions
>>> A0 = A[0]  # First partition
>>> A1 = A[1]  # Second partition
>>>
>>> # Iterate over partitions
>>> for partition in A:
>>>     x = partition.solve(b_local)
>>>
>>> # Properties
>>> print(A.num_partitions)  # 4
>>> print(A.shape)           # Global shape
>>> print(len(A))            # 4
>>>
>>> # Move to CUDA
>>> A_cuda = A.cuda()
>>>
>>> # Local halo exchange (for testing)
>>> x_list = [torch.zeros(A[i].num_local) for i in range(4)]
>>> A.halo_exchange_local(x_list)

classmethod from_sparse_local(local_tensor: SparseTensor, mesh: Any, partition: Partition, *, axis: int = 0, global_shape: Tuple[int, int] | None = None) → DSparseTensor

Wrap a per-rank SparseTensor chunk (already in local coords) plus its Partition as a DSparseTensor.

Use together with SparseTensor.extract_partition():

partition = compute_partition(...)
local_tensor = A_global.extract_partition(partition)
D = DSparseTensor.from_sparse_local(
    local_tensor, mesh, partition,
    global_shape=A_global.shape,
)
y_dt = D @ x_dt              # halo exchange + local SpMV

The partition is stamped onto _spec.placement.partition so the placement is the single source of truth for the irregular shard map.

Parameters:

• local_tensor (SparseTensor) – This rank’s local subdomain (size (num_local, num_local), COO in local coordinates). Usually built by SparseTensor.extract_partition().

• mesh (DeviceMesh) – The PyTorch device mesh.

• partition (Partition) – Irregular partition map for this rank (owned_nodes / halo_nodes / neighbor_partitions etc).

• axis (int) – Sparse axis being sharded (default 0 = rows).

• global_shape (Tuple[int, int], optional) – Global matrix shape. If omitted, inferred from partition.local_to_global.numel() – only valid for square matrices.

classmethod partition_batch(A: SparseTensor, mesh: Any, *, axis: int = 0) → DSparseTensor

Batch-shard a batched SparseTensor across mesh.

Every rank gets the same row/col indices; only the values tensor is sliced along A.batch_shape[axis]. No halo exchange, no cross-rank comm in matvec.

Requires A.is_batched and axis < len(A.batch_shape).

classmethod partition(A: SparseTensor, mesh: Any, *, partition_method: str = 'simple', coords: Tensor | None = None, verbose: bool = False) → DSparseTensor

One-shot constructor: take a global SparseTensor + DeviceMesh, partition rows across the mesh, return a ready-to-use distributed tensor with RowPartitioned placement.

\[\{0,\dots,N-1\} = \bigsqcup_{r=0}^{P-1} V_r,\qquad \text{halo}(r) = \{ j \notin V_r : \exists\, i\in V_r,\, A_{ij}\neq 0 \}\]

i.e. rows are split into P disjoint owned sets V_r (one per rank), and each rank additionally tracks the off-rank columns its rows touch (the halo) for matvec communication.

Algorithm

Compute a balanced row partition, then build each rank’s local sub-matrix + halo / send-recv index plan. "metis" / "rcb" minimize edge cut (halo volume); "simple" / "slicing" are contiguous splits. Determinism is enforced by computing labels on rank 0 and broadcasting:

ids   = resolve_partition_ids(A, P, method)   # rank 0, broadcast
part  = build_partition(A, ids, rank)         # owned + halo + plan
local = A.extract_partition(part)             # this rank's block
return DSparseTensor(local, mesh, RowPartitioned)

Complexity

Time \(O(nnz)\) for "simple"/"slicing"; graph partitioners (METIS) are \(O(nnz)\)-ish near-linear in practice. Space \(O(N + nnz/P)\) per rank (owned rows + halo).

Equivalent to:

local = A.partition_for_rank(rank, world_size,
                              partition_method=partition_method,
                              coords=coords)
D = DSparseTensor.from_local(local, mesh,
                              placement=RowPartitioned())

but in one line. This is the recommended way to build a distributed sparse tensor from a global SparseTensor for both unit tests and small-to-medium production runs (where every rank can afford to hold the global A briefly).

For memory-tight scenarios where only rank 0 should ever materialise the global matrix, use from_global_distributed() (which broadcasts only the partition IDs from rank 0) and chain from_local() manually.

param A:

Global sparse matrix; every rank should hold an identical copy at the time of the call.

type A:

SparseTensor

param mesh:

Target device mesh. mesh.size() becomes the world size and dist.get_rank() picks this rank’s chunk.

type mesh:

DeviceMesh

param partition_method:

Partitioning algorithm passed through to SparseTensor.partition_for_rank(): "simple" / "metis" / "rcb" / "slicing".

type partition_method:

str

param coords:

Node coordinates for geometric partitioning (RCB/slicing).

type coords:

torch.Tensor, optional

param verbose:

Print partition info on each rank.

type verbose:

bool

property spec: DSparseSpec | None

The DSparseSpec for this tensor (placement + mesh + global shape), or None if this instance was built via the legacy single-process simulator constructor.

scatter(global_vec: Tensor) → DTensor

Convenience: extract this rank’s owned slice from a global vector and wrap as a DTensor[Shard(0)].

Common usage:

b_dt = D.scatter(b_global)        # build distributed RHS
x_dt = solve(D, b_dt)             # distributed solve
r_dt = b_dt - D @ x_dt            # distributed residual

global_vec is a 1-D torch.Tensor of size global_shape[0]. Every rank should hold the same copy (typical in tests; in production the caller loads on rank 0 and broadcasts).

full_tensor() → SparseTensor

Materialise the full global tensor on every rank.

Mirrors DTensor.full_tensor(). For SparseShard(axis=0) we drop halo rows, translate indices to global, and allgather the COO triples. For BatchShard we allgather the per-rank values slices along the sharded batch axis.

classmethod from_global_distributed(values: Tensor, row_indices: Tensor, col_indices: Tensor, shape: Tuple[int, int], rank: int, world_size: int, mesh: Any = None, coords: Tensor | None = None, partition_method: str = 'auto', device: str | device | None = None, verbose: bool = True) → DSparseTensor

Create local partition in a distributed-safe manner.

This method ensures that all ranks compute the same partition assignment by having rank 0 compute the partition IDs and broadcasting to all ranks.

Parameters:

• values (torch.Tensor) – Global non-zero values [nnz]

• row_indices (torch.Tensor) – Global row indices [nnz]

• col_indices (torch.Tensor) – Global column indices [nnz]

• shape (Tuple[int, int]) – Global matrix shape (M, N)

• rank (int) – Current process rank

• world_size (int) – Total number of processes

• coords (torch.Tensor, optional) – Node coordinates for geometric partitioning [num_nodes, dim]

• partition_method (str) – Partitioning method: ‘metis’, ‘rcb’, ‘slicing’, ‘simple’

• device (str or torch.device, optional) – Target device

• verbose (bool) – Whether to print partition info

Returns:

This rank’s row-sharded distributed tensor.

Return type:

DSparseTensor

Example

>>> import torch.distributed as dist
>>>
>>> # In each process:
>>> rank = dist.get_rank()
>>> world_size = dist.get_world_size()
>>>
>>> local_matrix = DSparseTensor.from_global_distributed(
...     val, row, col, shape,
...     rank=rank, world_size=world_size
... )

save(directory: Any, rank: int | None = None, verbose: bool = False) → None

Persist this rank’s shard to directory. Convenience for torch_sla.io.save_dsparse(self, directory)().

classmethod load(directory: Any, mesh: Any = None, rank: int | None = None, target_world_size: int | None = None, device: str | device = 'cpu') → DSparseTensor

Reconstruct a DSparseTensor from a directory previously written by save() / save_dsparse() / save_sparse_sharded().

Pass target_world_size=1 (or call from a single process with no live torch.distributed group) to gather all shards into one trivial mesh=None DSparseTensor – useful for offline inspection of a sharded archive. If stored_num_partitions != target_world_size and the target is not 1, raises NotImplementedError (true cross-world-size repartition is deferred to a future redistribute()).

property shape: Tuple[int, int]

Global matrix shape.

property num_partitions: int

Number of partitions.

property device: device

Device of the matrix data.

property dtype: dtype

Data type of matrix values.

property nnz: int

Local nnz on this rank. Use global_nnz() for the sum.

global_nnz() → int

Sum of nnz across all ranks. Collective; cached.

property ndim: int

property sparse_shape: Tuple[int, int]

property sparse_dim: Tuple[int, int]

property batch_shape: Tuple[int, ...]

property block_shape: Tuple[int, ...]

property batch_size: int

property is_batched: bool

property is_block: bool

property is_cuda: bool

property is_square: bool

property values: Tensor

property row_indices: Tensor

property col_indices: Tensor

to(device: str | device | None = None, dtype: dtype | None = None) → DSparseTensor

cuda(device: int | None = None) → DSparseTensor

cpu() → DSparseTensor

float() → DSparseTensor

double() → DSparseTensor

half() → DSparseTensor

sum(axis: int | Tuple[int, ...] | None = None, keepdim: bool = False) → Tensor

Sum stored values. axis in {None, 0, 1}.

mean(axis: int | Tuple[int, ...] | None = None) → Tensor

Mean over stored values (implicit zeros excluded).

prod() → Tensor

max() → Tensor

min() → Tensor

norm(ord: Any = 'fro') → Tensor

'fro' / 1 / inf. 2 requires full_tensor().norm(2).

abs() → DSparseTensor

sqrt() → DSparseTensor

square() → DSparseTensor

exp() → DSparseTensor

log() → DSparseTensor

conj() → DSparseTensor

is_symmetric(atol: float = 1e-08, rtol: float = 1e-05) → bool

is_hermitian(atol: float = 1e-08, rtol: float = 1e-05) → bool

is_positive_definite() → bool

detect_matrix_type() → str

det() → Tensor

Determinant via full_tensor() + single-process LU. Warns.

lu()

LU factorisation via full_tensor() + single-process LU. Warns.

svd(k: int = 6)

Truncated SVD via full_tensor() + single-process. Warns.

condition_number(ord: int = 2) → Tensor

Condition number via full_tensor() + single-process. Warns.

logdet(**kwargs) → Tensor

Distributed log-determinant via Hutchinson + Lanczos.

When method='hutchinson' (the default for SPD), no gather happens – the trace estimator only needs A @ z which routes through _shard_matvec. Falls back to full_tensor() + single-process LU for non-SPD or explicit method='lu'.

See torch_sla.det for the full DetConfig knobs.

T() → DSparseTensor

Transpose. Allgathers, transposes, repartitions on same mesh.

H() → DSparseTensor

eigsh(k: int = 6, which: str = 'LM', maxiter: int = 200, tol: float = 1e-08, return_eigenvectors: bool = True, sigma: float | None = None, verbose: bool = False)

Distributed LOBPCG (SparseShard) or per-batch eigsh (BatchShard).

BatchShard returns (eigenvalues, eigenvectors) whose first axis is the batch axis – each rank runs SparseTensor.eigsh on its local batch slice, no inter-rank comm.

solve_batch_shard(b: Tensor, **kwargs) → Tensor

Per-batch solve under BatchShard. Each rank slices b to its own batch range and reuses SparseTensor.solve_batch() (same-pattern batched solve) on its local values stack. Returns this rank’s batch slice of the solution; allgather via full_tensor()-style code if you need it globally. Zero inter-rank communication.

__matmul__(x: Tensor | DTensor) → Tensor | DTensor

Distributed matrix–vector product D @ x.

\[y = A x,\qquad y_r = A_{r,V_r} x_{V_r} + A_{r,\text{halo}(r)} x_{\text{halo}(r)}\]

Each rank computes its owned slice y_r of the global product from its local rows; columns it does not own are fetched from neighbours via a halo exchange before the local SpMV.

Algorithm

x_halo = halo_exchange(x)        # one sparse all-to-all (neighbours)
x_loc  = concat(x_owned, x_halo)
y_r    = local_spmv(A_local, x_loc)   # O(local nnz)

Complexity

Time \(O(nnz / P)\) compute per rank plus one halo exchange whose volume is the edge cut (small for a good partition); space \(O(N/P + |\text{halo}|)\). Differentiable: backward is a halo exchange of the adjoint vector plus a transposed local SpMV (\(O(1)\) extra graph nodes).

See distributed_matvec.matmul_spec() / distributed_matvec.matmul_batch_shard().

solve_distributed_shard(b: Any, *, method: Any = None, preconditioner: Any = None, atol: Any = None, rtol: Any = None, maxiter: Any = None, restart: Any = None, verbose: Any = None) → Any

Distributed Krylov solve entirely in Shard(0) space.

Requires this DSparseTensor to carry a real spec (built via from_local()). The right-hand side b may be a DTensor[Shard(0)] (most common) or a raw torch.Tensor sized num_owned for the calling rank; the return value mirrors the input’s wrapper.

Methods (all live in Shard(0) space):

• "cg" Saad §6.7 conjugate gradient – SPD systems

• "bicgstab" Saad §7.4 BiCGStab – non-symmetric, no restart

• "gmres" Saad §6.5 restarted GMRES(m) – general

• "fgmres" Saad §9.4 flexible GMRES(m) – variable preconditioner

• "minres" Paige-Saunders MINRES – symmetric indefinite

Inner products go through dist.all_reduce (sum), matvec through halo_exchange – no rank ever sees a global vector.

SolverConfig integration

Every kwarg defaults to None, meaning “look at the active SolverConfig scope on this thread, then fall back to the hard-coded default”. The precedence chain matches solve() – explicit kwarg → innermost scope → outer scopes (LIFO) → hard-coded default.

>>> with SolverConfig(method="bicgstab", atol=1e-8):
...     x = D.solve_distributed_shard(b)        # picks BiCGStab + 1e-8
...     y = D.solve_distributed_shard(b, atol=1e-12)  # kwarg wins

nonlinear_solve_distributed_shard(residual_fn: Any, u0: Any, *, jac_diag_fn: Any, tol: Any = 1e-10, atol: Any = 1e-12, max_iter: Any = 50, line_search: bool = True, lin_maxiter: int = 1000, verbose: bool = False, adjoint_dLdu: Any = None) → Any

Distributed Newton solve for F(u) = 0 in Shard(0) space.

F must have the structure F(u) = A u + g(u) with A the distributed operator (this DSparseTensor) and g a pointwise (diagonal) nonlinearity. The Newton step solves the Jacobian system J du = -F, J v = A v + d * v, via distributed GMRES; the diagonal shift d = g'(u) is supplied by jac_diag_fn.

Parameters:

• residual_fn (Callable) – residual_fn(u_owned, D) -> F_owned using distributed ops.

• u0 (DTensor[Shard(0)] | torch.Tensor) – Initial guess (owned slice or wrapped).

• jac_diag_fn (Callable) – jac_diag_fn(u_owned, D) -> d_owned; Jacobian diagonal shift.

• adjoint_dLdu (DTensor | torch.Tensor, optional) – If given, also solve the IFT adjoint Jᵀ λ = dL/du at the converged u and return (u, λ) instead of just u.

• -> (Returns the solution mirroring the input wrapper (DTensor in)

• ``(u (DTensor out). With adjoint_dLdu set returns) –

• λ)``.

connected_components_distributed_shard()

Distributed connected components on the VertexShard graph.

Treats the matrix as an undirected adjacency and returns (labels_owned, num_components) for this rank’s owned slice; num_components and the labelling agree with the single-process / scipy result. See torch_sla.distributed.graph.connected_components_shard().

solve(b: Any, **kwargs) → Any

D.solve(b) – distributed Krylov solve in Shard(0) space.

\[A x = b \quad\Longrightarrow\quad x = A^{-1} b\]

Solves the global system with no rank ever materialising a global vector: every Krylov vector stays the owned slice (size num_owned).

Algorithm

Distributed CG / BiCGStab / GMRES / MINRES. Sparse matvec uses a halo exchange (D @ p); the two scalar inner products per iteration are summed across ranks with dist.all_reduce:

r = b - D @ x;  p = r
repeat until ||r|| <= atol:        # ||.|| via all_reduce
    Ap    = D @ p                  # halo_exchange + local SpMV
    alpha = all_reduce(<r,r>) / all_reduce(<p,Ap>)
    x += alpha*p;  r -= alpha*Ap;  update p

Complexity

Time \(O(m \cdot nnz / P)\) compute per rank over m iterations, plus m halo exchanges and O(m) all-reduces; space \(O((n + nnz)/P)\) per rank.

Sugar for solve_distributed_shard(); mirrors the top-level torch_sla.solve(). b is a DTensor[Shard(0)] (most common) or a raw owned-slice torch.Tensor; the return value matches the input’s wrapper. Method / preconditioner / tolerances resolve through the active SolverConfig scope unless passed explicitly (see solve_distributed_shard()).

solve_batch(b: Tensor, **kwargs) → Tensor

D.solve_batch(b) – per-batch solve under BatchShard.

Sugar for solve_batch_shard(). Returns this rank’s batch slice of the solution; zero inter-rank communication.

nonlinear_solve(residual_fn: Any, u0: Any, *, jac_diag_fn: Any, **kwargs) → Any

D.nonlinear_solve(residual, u0, jac_diag_fn=...) – distributed Newton solve for \(F(u) = 0\) in Shard(0) space.

\[F(u^\*) = 0,\qquad u_{n+1} = u_n - J(u_n)^{-1} F(u_n),\quad J=\partial F/\partial u\]

Algorithm

Newton-Raphson where every step’s Jacobian solve is a distributed Krylov solve (solve()); residual and Jacobian-diagonal are evaluated rank-locally:

u = u0
for it in 1..max_iter:
    F = residual(u)              # local, halo-exchanged
    if ||F|| <= tol: break       # ||.|| via all_reduce
    du = D_J.solve(-F)           # distributed Krylov (inner)
    u += du

Complexity

Time \(O(m \cdot \text{solve})\) for m Newton steps, each solve a distributed Krylov solve; space \(O((n+nnz)/P)\) per rank. Gradients (when requested) use the IFT adjoint – one extra distributed solve, \(O(1)\) graph nodes.

Sugar for nonlinear_solve_distributed_shard(); mirrors the single-process SparseTensor.nonlinear_solve(). Returns the solution mirroring the input wrapper, or (u, λ) when an adjoint_dLdu kwarg requests the IFT adjoint.

connected_components() → Tuple[Tensor, int]

D.connected_components() – distributed connected components.

\[G(A)=(V,E),\quad E=\{(i,j):A_{ij}\neq 0\} \;\longrightarrow\; \text{maximal connected subsets of } V\]

Algorithm

Distributed FastSV / Shiloach-Vishkin label propagation: every rank hooks its owned edges to the minimum neighbour label, exchanges boundary labels across ranks, and path-compresses, iterating to a global fixed point in \(O(\log N)\) rounds:

label[i] = i  (owned)
repeat until global labels stable:
    hook owned edges (scatter_reduce amin)
    halo_exchange(boundary labels)        # cross-rank
    compress_to_roots(label)

Complexity

Time \(O((nnz/P)\log N)\) compute per rank plus \(O(\log N)\) halo exchanges; space \(O(N/P)\). Non-differentiable (discrete labels).

Sugar for torch_sla.distributed.graph.connected_components_shard(); mirrors the single-process SparseTensor.connected_components().

returns:

• labels (torch.Tensor) – This rank’s owned-slice of the contiguous component labelling (0..n_components-1, identical to the scipy labelling).

• n_components (int) – Global component count (identical on every rank).

lsqr(b: Tensor, *, atol: float = 1e-08, btol: float = 1e-08, maxiter: int = 10000, damp: float = 0.0, conlim: float = 100000000.0, verbose: bool = False) → Tensor

D.lsqr(b) – distributed LSQR least-squares solve.

Sugar for torch_sla.distributed.solve.lsqr_shard(); the distributed counterpart of single-process spsolve(method='lsqr') (defaults match its lsqr_solve). b is this rank’s owned-slice torch.Tensor; returns the owned slice of argmin ||Ax - b|| (damped variant when damp > 0).

lsmr(b: Tensor, *, atol: float = 1e-08, btol: float = 1e-08, maxiter: int = 10000, damp: float = 0.0, conlim: float = 100000000.0, verbose: bool = False) → Tensor

D.lsmr(b) – distributed LSMR least-squares solve.

Sugar for torch_sla.distributed.solve.lsmr_shard(); the distributed counterpart of single-process spsolve(method='lsmr') (defaults match its lsmr_solve). b is this rank’s owned-slice torch.Tensor; returns the owned slice of the least-squares solution (more robust than LSQR when ill-conditioned).

Partition

Dataclass representing a single partition/subdomain for distributed computing.

class torch_sla.Partition(partition_id: int, local_nodes: Tensor, owned_nodes: Tensor, halo_nodes: Tensor, neighbor_partitions: List[int], send_indices: Dict[int, Tensor], recv_indices: Dict[int, Tensor], global_to_local: Tensor, local_to_global: Tensor)

One rank’s view of the global graph after partitioning.

partition_id

Type:

int

local_nodes

Global indices of every node visible to this rank (owned + halo).

Type:

torch.Tensor

owned_nodes

Subset actually owned – the rank’s slice of the row partition.

Type:

torch.Tensor

halo_nodes

Ghost nodes pulled in from neighbours each halo exchange.

Type:

torch.Tensor

neighbor_partitions

IDs of partitions this rank exchanges with.

Type:

List[int]

send_indices

For each neighbour, which of OUR owned nodes we send them.

Type:

Dict[int, torch.Tensor]

recv_indices

For each neighbour, which positions in the local layout the received data lands in.

Type:

Dict[int, torch.Tensor]

global_to_local

Global → local index map; -1 (or out-of-bounds) for nodes not in this partition.

Type:

torch.Tensor

local_to_global

Inverse map: local index → global node id.

Type:

torch.Tensor

partition_id: int

local_nodes: Tensor

owned_nodes: Tensor

halo_nodes: Tensor

neighbor_partitions: List[int]

send_indices: Dict[int, Tensor]

recv_indices: Dict[int, Tensor]

global_to_local: Tensor

local_to_global: Tensor

to(device: device | str) → Partition

Return a copy with every tensor field moved to device.

Index dtypes (int64) are preserved. Used by DSparseTensor.to() so halo exchange tensors live alongside the local matrix data.

---

Linear Solve Functions

spsolve

torch_sla.spsolve(val: Tensor, row: Tensor, col: Tensor, shape: Tuple[int, int], b: Tensor, backend: Literal['scipy', 'pytorch', 'cudss', 'pyamg', 'amgx', 'strumpack', 'auto'] = 'auto', method: Literal['auto', 'lu', 'umfpack', 'cholesky', 'ldlt', 'cg', 'cgs', 'bicgstab', 'gmres', 'lgmres', 'minres', 'qmr', 'lsqr', 'lsmr', 'amg', 'ruge_stuben', 'smoothed_aggregation', 'sa'] = 'auto', atol: float = 1e-10, maxiter: int = 10000, tol: float = 1e-12, matrix_type: str = 'general', is_symmetric: bool = False, is_spd: bool = False, preconditioner: str = 'jacobi', mixed_precision: bool = False, verbose: bool = False) → Tensor

Solve the Sparse Linear Equation Ax = b with gradient support.

Supports multiple backends for CPU and CUDA tensors.

Parameters:

• val (torch.Tensor) – [nnz] Non-zero values of sparse matrix A in COO format

• row (torch.Tensor) – [nnz] Row indices

• col (torch.Tensor) – [nnz] Column indices

• shape (Tuple[int, int]) – (m, n) Shape of sparse matrix A

• b (torch.Tensor) – [m] or [m, K] Right-hand side vector (or matrix for multiple RHS)

• backend (str, optional) – Backend to use: - ‘auto’: Auto-select based on device and problem size (default) - ‘scipy’: SciPy (CPU only, uses LU/UMFPACK) - ‘pytorch’: PyTorch-native (CPU / CUDA / ROCm, iterative) - best for large problems - ‘strumpack’: STRUMPACK multifrontal direct (CPU / CUDA / ROCm) - ‘cudss’: NVIDIA cuDSS (CUDA only, direct)

• method (str, optional) –

  Solver method. Available methods depend on backend: - ‘auto’: Auto-select based on matrix properties - ‘lu’: LU factorization (scipy, strumpack, cudss) - ‘umfpack’: UMFPACK direct solver (CPU only, via scipy/scikit-umfpack;

  there is no GPU/ROCm UMFPACK – for portable direct solves use ‘strumpack’)

  • ’cholesky’, ‘ldlt’: Direct solvers (cudss)

  • ’cg’, ‘cgs’, ‘bicgstab’, ‘gmres’: Iterative solvers

• atol (float, optional) – Absolute tolerance for iterative solvers, by default 1e-10

• maxiter (int, optional) – Maximum iterations for iterative solvers, by default 10000

• tol (float, optional) – Tolerance for direct solvers, by default 1e-12

• matrix_type (str, optional) – Matrix type for cuDSS: ‘general’, ‘symmetric’, ‘spd’, by default “general”

• is_symmetric (bool, optional) – Hint that matrix is symmetric (for auto method selection)

• is_spd (bool, optional) – Hint that matrix is symmetric positive definite

Returns:

[n] or [n, K] Solution vector (or matrix for multiple RHS)

Return type:

torch.Tensor

Examples

>>> import torch
>>> from torch_sla import spsolve
>>>
>>> # Create a simple SPD matrix
>>> val = torch.tensor([4.0, -1.0, -1.0, 4.0, -1.0, -1.0, 4.0], dtype=torch.float64)
>>> row = torch.tensor([0, 0, 1, 1, 1, 2, 2], dtype=torch.int64)
>>> col = torch.tensor([0, 1, 0, 1, 2, 1, 2], dtype=torch.int64)
>>> shape = (3, 3)
>>> b = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64)
>>>
>>> # Auto-select backend and method
>>> x = spsolve(val, row, col, shape, b)
>>>
>>> # Specify backend and method
>>> x = spsolve(val, row, col, shape, b, backend='scipy', method='lu')
>>>
>>> # On CUDA
>>> val_cuda = val.cuda()
>>> row_cuda = row.cuda()
>>> col_cuda = col.cuda()
>>> b_cuda = b.cuda()
>>> x_cuda = spsolve(val_cuda, row_cuda, col_cuda, shape, b_cuda, backend='cudss', method='lu')

spsolve_coo

torch_sla.spsolve_coo(A: Tensor, b: Tensor, **kwargs) → Tensor

Solve Ax = b where A is a sparse COO tensor

Parameters:

• A (torch.Tensor) – Sparse COO tensor representing the matrix

• b (torch.Tensor) – Right-hand side vector

• **kwargs – Additional arguments passed to spsolve()

Returns:

Solution vector x

Return type:

torch.Tensor

spsolve_csr

torch_sla.spsolve_csr(A: Tensor, b: Tensor, **kwargs) → Tensor

Solve Ax = b where A is a sparse CSR tensor

Parameters:

• A (torch.Tensor) – Sparse CSR tensor representing the matrix

• b (torch.Tensor) – Right-hand side vector

• **kwargs – Additional arguments passed to spsolve()

Returns:

Solution vector x

Return type:

torch.Tensor

---

Batch Solve Functions

spsolve_batch_same_layout

torch_sla.spsolve_batch_same_layout(val_batch: Tensor, row: Tensor, col: Tensor, shape: Tuple[int, int], b_batch: Tensor, method: Literal['cg', 'bicgstab', 'gmres', 'minres', 'lsqr', 'lsmr', 'strumpack', 'cudss', 'cudss_lu', 'cudss_cholesky', 'cudss_ldlt'] = 'bicgstab', atol: float = 1e-10, maxiter: int = 10000) → Tensor

Batch solve sparse linear systems with the SAME sparsity pattern.

Note

method='cg' runs a truly batched solver (one scatter matvec over the shared pattern, all systems iterate together) – this is the fast path on GPU / ROCm. Other methods (bicgstab/gmres/minres/lsqr/lsmr, strumpack, cudss_*) currently solve each system in a Python loop.

All matrices A_i share the same (row, col) structure but have different values. This is efficient when the sparsity pattern is fixed (e.g., FEM with fixed mesh).

Solves: A_i @ x_i = b_i for i = 0, 1, …, batch_size-1

Parameters:

• val_batch (torch.Tensor) – [batch_size, nnz] Non-zero values for each matrix

• row (torch.Tensor) – [nnz] Row indices (shared across batch)

• col (torch.Tensor) – [nnz] Column indices (shared across batch)

• shape (Tuple[int, int]) – (m, n) Shape of each sparse matrix

• b_batch (torch.Tensor) – [batch_size, m] Right-hand side vectors

• method (str) – Solver method (same options as spsolve)

• atol (float) – Absolute tolerance for iterative solvers

• maxiter (int) – Maximum iterations for iterative solvers

Returns:

[batch_size, n] Solution vectors

Return type:

torch.Tensor

Example

>>> import torch
>>> from torch_sla import spsolve_batch_same_layout
>>>
>>> batch_size = 10
>>> n = 100
>>> nnz = 500
>>>
>>> # Same sparsity pattern, different values
>>> row = torch.randint(0, n, (nnz,))
>>> col = torch.randint(0, n, (nnz,))
>>> val_batch = torch.randn(batch_size, nnz, dtype=torch.float64)
>>> b_batch = torch.randn(batch_size, n, dtype=torch.float64)
>>>
>>> x_batch = spsolve_batch_same_layout(val_batch, row, col, (n, n), b_batch)

spsolve_batch_different_layout

torch_sla.spsolve_batch_different_layout(matrices: List[Tuple[Tensor, Tensor, Tensor, Tuple[int, int]]], b_list: List[Tensor], method: Literal['cg', 'bicgstab', 'gmres', 'minres', 'lsqr', 'lsmr', 'strumpack', 'cudss', 'cudss_lu', 'cudss_cholesky', 'cudss_ldlt'] = 'bicgstab', atol: float = 1e-10, maxiter: int = 10000) → List[Tensor]

Batch solve sparse linear systems with DIFFERENT sparsity patterns.

Deprecated since version Use: SparseTensorList.solve() instead for a more Pythonic interface:

>>> matrices = SparseTensorList([A1, A2, A3])
>>> x_list = matrices.solve([b1, b2, b3])

Each matrix can have a different structure. This is useful when dealing with heterogeneous problems or adaptive mesh refinement.

Parameters:

• matrices (List[Tuple[val, row, col, shape]]) – List of sparse matrices, each as (values, row_indices, col_indices, shape)

• b_list (List[torch.Tensor]) – List of right-hand side vectors

• method (str) – Solver method (same options as spsolve)

• atol (float) – Absolute tolerance for iterative solvers

• maxiter (int) – Maximum iterations for iterative solvers

Returns:

List of solution vectors

Return type:

List[torch.Tensor]

Example

>>> import torch
>>> from torch_sla import spsolve_batch_different_layout
>>>
>>> # Different matrices with different sizes/patterns
>>> matrices = []
>>> b_list = []
>>> for n in [50, 100, 150]:
...     nnz = n * 5
...     val = torch.randn(nnz, dtype=torch.float64)
...     row = torch.randint(0, n, (nnz,))
...     col = torch.randint(0, n, (nnz,))
...     matrices.append((val, row, col, (n, n)))
...     b_list.append(torch.randn(n, dtype=torch.float64))
>>>
>>> x_list = spsolve_batch_different_layout(matrices, b_list)

ParallelBatchSolver

class torch_sla.ParallelBatchSolver(row: Tensor, col: Tensor, shape: Tuple[int, int], method: Literal['cg', 'bicgstab', 'gmres', 'minres', 'lsqr', 'lsmr', 'strumpack', 'cudss', 'cudss_lu', 'cudss_cholesky', 'cudss_ldlt'] = 'bicgstab', device: str | None = None)

Bases: object

High-performance parallel batch solver.

This class pre-analyzes the sparsity pattern and caches factorization information for repeated solves with the same structure.

Example

>>> solver = ParallelBatchSolver(row, col, shape, method='cudss_lu')
>>>
>>> # Solve multiple batches efficiently
>>> for val_batch, b_batch in data_loader:
...     x_batch = solver.solve(val_batch, b_batch)

solve(val_batch: Tensor, b_batch: Tensor, atol: float = 1e-10, maxiter: int = 10000) → Tensor

Solve batch of linear systems.

Parameters:

• val_batch (torch.Tensor) – [batch_size, nnz] Matrix values

• b_batch (torch.Tensor) – [batch_size, m] Right-hand side vectors

• atol (float) – Tolerance for iterative solvers

• maxiter (int) – Maximum iterations

Returns:

[batch_size, n] Solution vectors

Return type:

torch.Tensor

---

Nonlinear Solve

nonlinear_solve

torch_sla.nonlinear_solve(residual_fn: Callable, u0: Tensor, *params, jacobian_fn: Callable | None = None, method: str = 'newton', tol: float = 1e-06, atol: float = 1e-10, max_iter: int = 50, line_search: bool = True, verbose: bool = False, linear_solver: str = 'pytorch', linear_method: str = 'cg') → Tensor

Solve nonlinear equation F(u, θ) = 0 with adjoint-based gradients.

Parameters:

• residual_fn – Function F(u, *params) -> residual tensor

• u0 – Initial guess for solution

• *params – Parameters θ (tensors with requires_grad=True for gradient computation)

• jacobian_fn – Optional function J(u, *params) -> (val, row, col, shape) Returns sparse Jacobian in COO format. If None, uses autograd.

• method – Nonlinear solver method - ‘newton’: Newton-Raphson with optional line search (default) - ‘picard’: Fixed-point iteration - ‘anderson’: Anderson acceleration

• tol – Relative convergence tolerance

• atol – Absolute convergence tolerance

• max_iter – Maximum number of nonlinear iterations

• line_search – Use Armijo line search for Newton (default: True)

• verbose – Print convergence information

• linear_solver – Backend for linear solves (‘pytorch’, ‘scipy’, ‘cudss’)

• linear_method – Method for linear solves (‘cg’, ‘bicgstab’, ‘lu’)

Returns:

Solution tensor satisfying F(u, θ) ≈ 0

Return type:

u

Example

>>> def residual(u, A_val, b):
...     # Nonlinear: A(u) @ u - b where A depends on u
...     return torch.sparse.mm(A, u.unsqueeze(1)).squeeze() - b
...
>>> u0 = torch.zeros(n, requires_grad=False)
>>> A_val = torch.randn(nnz, requires_grad=True)
>>> b = torch.randn(n, requires_grad=True)
>>>
>>> u = nonlinear_solve(residual, u0, A_val, b, method='newton')
>>> loss = some_loss(u)
>>> loss.backward()  # Computes ∂L/∂A_val and ∂L/∂b via adjoint

adjoint_solve

torch_sla.adjoint_solve(residual_fn: Callable, u0: Tensor, *params, jacobian_fn: Callable | None = None, method: str = 'newton', tol: float = 1e-06, atol: float = 1e-10, max_iter: int = 50, line_search: bool = True, verbose: bool = False, linear_solver: str = 'pytorch', linear_method: str = 'cg') → Tensor

Solve nonlinear equation F(u, θ) = 0 with adjoint-based gradients.

Parameters:

• residual_fn – Function F(u, *params) -> residual tensor

• u0 – Initial guess for solution

• *params – Parameters θ (tensors with requires_grad=True for gradient computation)

• jacobian_fn – Optional function J(u, *params) -> (val, row, col, shape) Returns sparse Jacobian in COO format. If None, uses autograd.

• method – Nonlinear solver method - ‘newton’: Newton-Raphson with optional line search (default) - ‘picard’: Fixed-point iteration - ‘anderson’: Anderson acceleration

• tol – Relative convergence tolerance

• atol – Absolute convergence tolerance

• max_iter – Maximum number of nonlinear iterations

• line_search – Use Armijo line search for Newton (default: True)

• verbose – Print convergence information

• linear_solver – Backend for linear solves (‘pytorch’, ‘scipy’, ‘cudss’)

• linear_method – Method for linear solves (‘cg’, ‘bicgstab’, ‘lu’)

Returns:

Solution tensor satisfying F(u, θ) ≈ 0

Return type:

u

Example

>>> def residual(u, A_val, b):
...     # Nonlinear: A(u) @ u - b where A depends on u
...     return torch.sparse.mm(A, u.unsqueeze(1)).squeeze() - b
...
>>> u0 = torch.zeros(n, requires_grad=False)
>>> A_val = torch.randn(nnz, requires_grad=True)
>>> b = torch.randn(n, requires_grad=True)
>>>
>>> u = nonlinear_solve(residual, u0, A_val, b, method='newton')
>>> loss = some_loss(u)
>>> loss.backward()  # Computes ∂L/∂A_val and ∂L/∂b via adjoint

NonlinearSolveAdjoint

class torch_sla.NonlinearSolveAdjoint(*args, **kwargs)

Bases: Function

Adjoint-based nonlinear solver with automatic differentiation.

Uses implicit differentiation to compute gradients without storing intermediate Jacobians. Memory-efficient for large-scale problems.

static forward(ctx, u0: Tensor, num_params: int, *args) → Tensor

Forward pass: solve F(u, θ) = 0 for u.

Parameters:

• u0 – Initial guess for solution

• num_params – Number of parameter tensors

• *args – First num_params elements are param tensors, last is config dict

Returns:

Solution satisfying F(u, θ) ≈ 0

Return type:

u

static backward(ctx, grad_u: Tensor)

Backward pass using adjoint method.

Computes ∂L/∂θ = -λᵀ · ∂F/∂θ where (∂F/∂u)ᵀ · λ = grad_u

Returns:

(grad_u0, grad_num_params, *grad_params, grad_config)

Return type:

Tuple of gradients

---

Persistence (I/O)

safetensors Format

save_sparse

torch_sla.save_sparse(tensor: SparseTensor, path: str | Path, metadata: Dict[str, str] | None = None) → None

Save a SparseTensor to safetensors format.

Parameters:

• tensor (SparseTensor) – The sparse tensor to save.

• path (str or Path) – Output file path (should end with .safetensors).

• metadata (dict, optional) – Additional metadata to store in the file.

Example

>>> A = SparseTensor(val, row, col, (100, 100))
>>> save_sparse(A, "matrix.safetensors")

load_sparse

torch_sla.load_sparse(path: str | Path, device: str | device = 'cpu') → SparseTensor

Load a SparseTensor from safetensors format.

Parameters:

• path (str or Path) – Input file path.

• device (str or torch.device) – Device to load tensors to.

Returns:

The loaded sparse tensor.

Return type:

SparseTensor

Example

>>> A = load_sparse("matrix.safetensors", device="cuda")

save_distributed

load_partition

load_metadata

torch_sla.load_metadata(directory: str | Path) → Dict

Load metadata from a distributed sparse tensor directory.

Parameters:

directory (str or Path) – Directory containing partitioned data.

Returns:

Metadata including shape, dtype, num_partitions, etc.

Return type:

dict

Example

>>> meta = load_metadata("matrix_dist")
>>> print(f"Shape: {meta['shape']}, Partitions: {meta['num_partitions']}")

load_sparse_as_partition

load_distributed_as_sparse

save_dsparse

torch_sla.save_dsparse(tensor: DSparseTensor, directory: str | Path, rank: int | None = None, verbose: bool = False) → None

Save this rank’s shard. Every rank calls; rank 0 also writes metadata.json.

load_dsparse

torch_sla.load_dsparse(directory: str | Path, mesh=None, rank: int | None = None, target_world_size: int | None = None, device: str | device = 'cpu') → DSparseTensor

Load a sharded archive into a DSparseTensor.

Three modes, picked from target_world_size (or, if omitted, a live process group, else 1):

• stored_N == target_world_size – fast path; this rank reads partition_{rank}.safetensors directly. Original save_dsparse() semantics.

• target_world_size == 1 – single-process gather: read every shard from disk, stitch the owned rows into one global SparseTensor, return as a mesh=None trivial DSparseTensor. Useful for inspection / single-node debug after sharded training.

• otherwise – explicit NotImplementedError. True re-partitioning on load (stored_N != target != 1) requires a cross-rank shuffle; deferred to a future redistribute() pass. Workaround: load with target_world_size=1 and call .full_tensor() (or save again at the new world size).

Matrix Market Format

save_mtx

torch_sla.save_mtx(tensor: SparseTensor, path: str | Path, comment: str = '', field: str = 'real', symmetry: str = 'general') → None

Save a SparseTensor to Matrix Market (.mtx) format.

Parameters:

• tensor (SparseTensor) – The sparse tensor to save.

• path (str or Path) – Output file path (should end with .mtx).

• comment (str, optional) – Comment to include in the header.

• field (str, optional) – Field type: ‘real’, ‘complex’, ‘integer’, or ‘pattern’. Default: ‘real’.

• symmetry (str, optional) – Symmetry type: ‘general’, ‘symmetric’, ‘skew-symmetric’, or ‘hermitian’. Default: ‘general’.

Example

>>> A = SparseTensor(val, row, col, (100, 100))
>>> save_mtx(A, "matrix.mtx")
>>> save_mtx(A, "matrix.mtx", symmetry="symmetric")

load_mtx

torch_sla.load_mtx(path: str | Path, dtype: dtype | None = None, device: str | device = 'cpu') → SparseTensor

Load a SparseTensor from Matrix Market (.mtx) format.

Parameters:

• path (str or Path) – Input file path.

• dtype (torch.dtype, optional) – Data type for values. If None, inferred from file.

• device (str or torch.device) – Device to load tensors to.

Returns:

The loaded sparse tensor.

Return type:

SparseTensor

Example

>>> A = load_mtx("matrix.mtx")
>>> A = load_mtx("matrix.mtx", dtype=torch.float32, device="cuda")

load_mtx_info

torch_sla.load_mtx_info(path: str | Path) → Dict

Read Matrix Market file header without loading data.

Parameters:

path (str or Path) – Input file path.

Returns:

Dictionary with keys: ‘shape’, ‘nnz’, ‘field’, ‘symmetry’.

Return type:

dict

Example

>>> info = load_mtx_info("matrix.mtx")
>>> print(f"Shape: {info['shape']}, NNZ: {info['nnz']}")

---

Partitioning Functions

partition_graph_metis

torch_sla.partition_graph_metis(row: Tensor, col: Tensor, num_nodes: int, num_parts: int) → Tensor

Partition a graph using METIS (via pymetis).

Falls back to partition_simple() if pymetis isn’t importable.

Returns:

partition_ids – Partition ID per node, shape (num_nodes,).

Return type:

torch.Tensor

partition_coordinates

torch_sla.partition_coordinates(coords: Tensor, num_parts: int, method: str = 'rcb') → Tensor

Geometric partitioning by node coordinates.

Three methods are supported:

• "rcb" – Recursive Coordinate Bisection (median-cut along longest axis). Standard CFD/FEM partitioner.

• "slicing" – Sort along longest axis, slice into equal chunks. Fast but worse quality than RCB.

• "hilbert" – Sort by Hilbert space-filling-curve index, slice into equal chunks. ~10-100x faster than METIS for PDE/mesh graphs where geometric locality correlates with sparse adjacency. 2-D and 3-D only.

Parameters:

• coords (torch.Tensor) – Node coordinates, shape (num_nodes, dim).

• num_parts (int) – Number of partitions (power of 2 recommended for RCB).

• method (str) – "rcb" / "slicing" / "hilbert".

partition_simple

torch_sla.partition_simple(num_nodes: int, num_parts: int) → Tensor

Contiguous 1-D partitioning: rank k owns the k-th equal slice of [0, num_nodes). Fast, no quality guarantees.

---

Backend Utilities

get_available_backends

torch_sla.get_available_backends() → List[str]

Get list of available backends

show_backends

torch_sla.show_backends() → None

Print a formatted status report of all backends.

Shows which backends are available on the current machine and gives installation hints for the ones that are not. Useful right after pip install torch-sla to verify the runtime environment.

Example

>>> import torch_sla
>>> torch_sla.show_backends()
torch-sla backend status (CUDA: available)
  scipy    [CPU]      available
  pytorch  [CPU/CUDA] available
  cudss    [CUDA]     not available — pip install torch-sla[cudss]

get_backend_methods

torch_sla.get_backend_methods(backend: str) → List[str]

Get list of methods supported by a backend

get_default_method

torch_sla.get_default_method(backend: str) → str

Get default method for a backend

select_backend

torch_sla.select_backend(device: device, n: int | None = None, dtype: dtype | None = None, prefer_direct: bool = True) → str

Auto-select the best backend based on device, problem size, and dtype.

Recommendations based on benchmark results: - CPU: scipy+lu (all sizes, fast + machine precision) - CUDA (DOF < 2M): cudss+cholesky (fast + high precision) - CUDA (DOF >= 2M): pytorch+cg (memory efficient, ~1e-6 precision)

Parameters:

• device (torch.device) – Target device (cpu or cuda)

• n (int, optional) – Problem size (DOF). If > CUDA_ITERATIVE_THRESHOLD, prefer iterative.

• dtype (torch.dtype, optional) – Data type.

• prefer_direct (bool) – If True, prefer direct solvers over iterative (when applicable)

Returns:

Backend name (‘scipy’, ‘pytorch’, or ‘cudss’)

Return type:

str

select_method

torch_sla.select_method(backend: str, is_symmetric: bool = False, is_spd: bool = False, prefer_direct: bool = True) → str

Auto-select the best method for a given backend and matrix properties.

Recommendations based on benchmark results: - scipy: lu (direct, best precision) or cg (iterative, for SPD) - cudss: cholesky (SPD, fastest) > ldlt (symmetric) > lu (general) - pytorch: cg (SPD) or bicgstab (general), both with Jacobi preconditioning

Parameters:

• backend (str) – Backend name

• is_symmetric (bool) – Whether the matrix is symmetric

• is_spd (bool) – Whether the matrix is symmetric positive definite

• prefer_direct (bool) – If True, prefer direct solvers

Returns:

Method name

Return type:

str

Backend Availability Checks

torch_sla.is_scipy_available() → bool

Check if SciPy backend is available

torch_sla.is_cudss_available() → bool

Check if cuDSS backend is available (via nvmath-python)

torch_sla.backends.is_strumpack_available() → bool

Check if the STRUMPACK backend is available (CPU / CUDA / ROCm).

Requires the optional torch-strumpack package, whose compiled STRUMPACK extension must load on this machine. Install a platform wheel:

pip install torch-strumpack        # cpu / cuda / rocm

This is torch-sla’s portable (incl. AMD ROCm) sparse-direct path.

---

Utility Functions

auto_select_method

torch_sla.auto_select_method(nnz: int, n: int, dtype: dtype, is_cuda: bool, is_spd: bool = False, memory_threshold: float = 0.8) → Tuple[str, str]

Automatically select the best backend and method.

Parameters:

• nnz (int) – Number of non-zero elements.

• n (int) – Matrix dimension.

• dtype (torch.dtype) – Data type of the matrix.

• is_cuda (bool) – Whether the matrix is on CUDA.

• is_spd (bool, optional) – Whether the matrix is symmetric positive definite. Default: False.

• memory_threshold (float, optional) – Fraction of GPU memory to use. Default: 0.8.

Returns:

(backend, method) tuple.

Return type:

Tuple[str, str]

estimate_direct_solver_memory

torch_sla.estimate_direct_solver_memory(nnz: int, n: int, dtype: dtype) → int

Estimate memory required for direct sparse solver.

Parameters:

• nnz (int) – Number of non-zero elements.

• n (int) – Matrix dimension.

• dtype (torch.dtype) – Data type of the matrix.

Returns:

Estimated memory in bytes.

Return type:

int

get_available_gpu_memory

torch_sla.get_available_gpu_memory() → int

Get available GPU memory in bytes.

Returns:

Available GPU memory in bytes, or 0 if CUDA is not available.

Return type:

int

---

Constants

BACKEND_METHODS

Dictionary mapping backend names to available solver methods.

BACKEND_METHODS = {
    'scipy': ['lu', 'umfpack', 'cg', 'bicgstab', 'gmres', 'lgmres', 'minres', 'qmr'],
    'pytorch': ['cg', 'bicgstab', 'gmres', 'minres', 'lsqr', 'lsmr'],  # device-agnostic (CPU/CUDA/ROCm)
    'cudss': ['lu', 'cholesky', 'ldlt'],                               # NVIDIA CUDA only
    'pyamg': ['amg', 'ruge_stuben', 'smoothed_aggregation', 'sa'],
    'amgx': ['amg', 'cg', 'pcg', 'bicgstab', 'pbicgstab', 'gmres', 'fgmres'],
    'strumpack': ['lu'],                                               # multifrontal direct (CPU/CUDA/ROCm)
}

DEFAULT_METHODS

Dictionary mapping backend names to their default solver methods.

DEFAULT_METHODS = {
    'scipy': 'lu',
    'pytorch': 'cg',
    'cudss': 'cholesky',
    'pyamg': 'ruge_stuben',
    'amgx': 'pbicgstab',
    'strumpack': 'lu',
}

Type Aliases

• BackendType: Literal type for backend names: 'scipy', 'pytorch', 'cudss', 'pyamg', 'amgx', 'strumpack', 'auto'

• MethodType: Literal type for solver methods: 'lu', 'umfpack', 'cg', 'cgs', 'bicgstab', 'gmres', 'minres', 'cholesky', 'ldlt', 'lsqr', 'lsmr'