API 参考¶

本节提供 torch-sla 的完整 API 文档。

核心类¶

SparseTensor¶

稀疏矩阵操作的主类。支持批量操作、自动微分和多种后端。

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

基类：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

参数:

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, ...]

示例

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.

参数:

A (torch.Tensor) -- Dense tensor with shape [...batch, M, N, ...block].
sparse_dim (Tuple[int, int], optional) -- Which dimensions are sparse. Default: (-2, -1).

返回:

Sparse representation of A.

返回类型:

SparseTensor

示例

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

参数:: A (torch.Tensor) -- PyTorch sparse COO or CSR tensor (2D only).
返回:: SparseTensor representation.
返回类型:: SparseTensor

示例

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

参数:

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

返回:

New SparseTensor on the target device/dtype.

返回类型:

SparseTensor

示例

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

参数:: device (int, optional) -- CUDA device index. Default: current device.
返回:: Tensor on CUDA.
返回类型:: SparseTensor

cpu() → SparseTensor[源代码]¶

Move tensor to CPU.

返回:: Tensor on CPU.
返回类型:: 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.

备注

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

示例

>>> 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)[源代码]¶

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

示例

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

示例

>>> # 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

示例

>>> # 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

示例

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

示例

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

备注

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.

示例

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

备注

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

示例

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

备注

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().

示例

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

示例

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

备注

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

示例

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

示例

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

示例

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

参数:

path (str or PathLike) -- Output file path (should end with .safetensors).
metadata (dict, optional) -- Additional metadata to store.

示例

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

参数:

path (str or PathLike) -- Input file path.
device (str or torch.device) -- Device to load tensors to.

返回:

The loaded sparse tensor.

返回类型:

SparseTensor

示例

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

SparseTensorList¶

不同稀疏模式的多个稀疏矩阵容器。适用于异构图的批量操作。

class torch_sla.SparseTensorList(tensors: List[SparseTensor])[源代码]¶

基类：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.

参数:: 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

示例

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

参数:: matrices (List[Tuple]) -- List of (values, row_indices, col_indices, shape) tuples.
返回:: List of SparseTensors.
返回类型:: SparseTensorList

示例

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

参数:: A_list (List[torch.Tensor]) -- List of PyTorch sparse COO tensors.
返回:: List of SparseTensors.
返回类型:: 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.

参数:: device (str or torch.device) -- Target device.
返回:: New list with tensors on target device.
返回类型:: 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.

参数:: 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.
返回:: If axis is None: List of scalar tensors (one per matrix). If axis is 0 or 1: List of 1D tensors.
返回类型:: List[torch.Tensor] or torch.Tensor

示例

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

参数:: axis (int, optional) -- Same as sum().
返回:: List of mean values/vectors.
返回类型:: 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.

参数:

b_list (List[torch.Tensor]) -- List of right-hand side vectors, one per matrix.
**kwargs -- Additional arguments passed to SparseTensor.solve().

返回:

List of solutions.

返回类型:

List[torch.Tensor]

示例

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

is_symmetric(**kwargs) → List[Tensor][源代码]¶

Check symmetry for all matrices.

参数:: **kwargs -- Arguments passed to SparseTensor.is_symmetric().
返回:: List of boolean tensors.
返回类型:: List[torch.Tensor]

is_positive_definite(**kwargs) → List[Tensor][源代码]¶

Check positive definiteness for all matrices.

参数:: **kwargs -- Arguments passed to SparseTensor.is_positive_definite().
返回:: List of boolean tensors.
返回类型:: List[torch.Tensor]

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

Compute norms for all matrices.

参数:: ord ({'fro', 1, 2}) -- Norm type.
返回:: List of norm values.
返回类型:: List[torch.Tensor]

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

Compute eigenvalues for all matrices.

参数:

k (int) -- Number of eigenvalues.
**kwargs -- Additional arguments.

返回:

List of (eigenvalues, eigenvectors) tuples.

返回类型:

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

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

Compute eigenvalues for symmetric matrices.

参数:

k (int) -- Number of eigenvalues.
**kwargs -- Additional arguments.

返回:

List of (eigenvalues, eigenvectors) tuples.

返回类型:

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

svd(k: int = 6) → List[Tuple[Tensor, Tensor, Tensor]][源代码]¶

Compute SVD for all matrices.

参数:: k (int) -- Number of singular values.
返回:: List of (U, S, Vt) tuples.
返回类型:: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]

condition_number(ord: int = 2) → List[Tensor][源代码]¶

Compute condition numbers for all matrices.

参数:: ord (int) -- Norm order.
返回:: List of condition numbers.
返回类型:: List[torch.Tensor]

det() → List[Tensor][源代码]¶

Compute determinants for all matrices.

返回:: List of determinant values.
返回类型:: List[torch.Tensor]

示例

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

参数:

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().

返回:

fig -- The figure object.

返回类型:

matplotlib.figure.Figure

示例

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

返回:: Block-diagonal matrix with shape (sum(M_i), sum(N_i)).
返回类型:: SparseTensor

备注

The resulting matrix has the structure:

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

示例

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

参数:

sparse (SparseTensor) -- Block-diagonal matrix to split.
sizes (List[Tuple[int, int]]) -- List of (rows, cols) for each block. Must sum to sparse.shape.

返回:

List of extracted blocks.

返回类型:

SparseTensorList

示例

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

返回:: List of (M, N) tuples.
返回类型:: 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 分解，用于相同矩阵的高效重复求解。

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

基类：object

LU factorization wrapper for efficient repeated solves.

Created by SparseTensor.lu().

参数:

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.

示例

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

参数:: b (torch.Tensor) -- Right-hand side vector.
返回:: Solution x.
返回类型:: torch.Tensor

分布式类¶

DSparseTensor¶

支持域分解的分布式稀疏张量。使用 halo 交换进行分区间通信。

class torch_sla.DSparseTensor[源代码]¶

基类：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.

参数:

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

示例

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

参数:

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.

参数:

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

返回:

This rank's row-sharded distributed tensor.

返回类型:

DSparseTensor

示例

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

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.

参数:

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

线性求解函数¶

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.

参数:

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

返回:

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

返回类型:

torch.Tensor

示例

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

参数:

A (torch.Tensor) -- Sparse COO tensor representing the matrix
b (torch.Tensor) -- Right-hand side vector
**kwargs -- Additional arguments passed to spsolve()

返回:

Solution vector x

返回类型:

torch.Tensor

spsolve_csr¶

torch_sla.spsolve_csr(A: Tensor, b: Tensor, **kwargs) → Tensor[源代码]¶

Solve Ax = b where A is a sparse CSR tensor

参数:

A (torch.Tensor) -- Sparse CSR tensor representing the matrix
b (torch.Tensor) -- Right-hand side vector
**kwargs -- Additional arguments passed to spsolve()

返回:

Solution vector x

返回类型:

torch.Tensor

批量求解函数¶

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.

备注

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

参数:

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

返回:

[batch_size, n] Solution vectors

返回类型:

torch.Tensor

示例

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

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

参数:

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

返回:

List of solution vectors

返回类型:

List[torch.Tensor]

示例

>>> 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)

非线性求解¶

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.

参数:

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')

返回:

Solution tensor satisfying F(u, θ) ≈ 0

返回类型:

示例

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

参数:

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')

返回:

Solution tensor satisfying F(u, θ) ≈ 0

返回类型:

示例

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

持久化 (I/O)¶

safetensors 格式¶

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

Save a SparseTensor to safetensors format.

参数:

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.

示例

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

torch_sla.load_sparse(path: str | Path, device: str | device = 'cpu') → SparseTensor[源代码]¶

Load a SparseTensor from safetensors format.

参数:

path (str or Path) -- Input file path.
device (str or torch.device) -- Device to load tensors to.

返回:

The loaded sparse tensor.

返回类型:

SparseTensor

示例

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

Matrix Market 格式¶

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.

参数:

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

示例

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

torch_sla.load_mtx(path: str | Path, dtype: dtype | None = None, device: str | device = 'cpu') → SparseTensor[源代码]¶

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

参数:

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.

返回:

The loaded sparse tensor.

返回类型:

SparseTensor

示例

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

后端工具¶

torch_sla.get_available_backends() → List[str][源代码]¶: Get list of available backends

torch_sla.get_backend_methods(backend: str) → List[str][源代码]¶: Get list of methods supported by a 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)

参数:

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)

返回:

Backend name ('scipy', 'pytorch', or 'cudss')

返回类型:

str

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

参数:

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

返回:

Method name

返回类型:

str

常量¶

BACKEND_METHODS¶

后端名称到可用求解方法的映射字典。

BACKEND_METHODS = {
    'scipy': ['lu', 'umfpack', 'cg', 'bicgstab', 'gmres', 'lgmres', 'minres', 'qmr'],
    'pytorch': ['cg', 'bicgstab', 'gmres', 'minres', 'lsqr', 'lsmr'],  # 设备无关 (CPU/CUDA/ROCm)
    'cudss': ['lu', 'cholesky', 'ldlt'],                               # 仅 NVIDIA CUDA
    'pyamg': ['amg', 'ruge_stuben', 'smoothed_aggregation', 'sa'],
    'amgx': ['amg', 'cg', 'pcg', 'bicgstab', 'pbicgstab', 'gmres', 'fgmres'],
    'strumpack': ['lu'],                                               # 多波前直接求解 (CPU/CUDA/ROCm)
}

DEFAULT_METHODS¶

后端名称到默认求解方法的映射字典。

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