Utility Functions

This module contains utility functions for sparse tensor operations.

Iterative Solvers

class torchsparsegradutils.utils.bicgstab.BICGSTABSettings(matvec_max, abstol, reltol, precon, logger)[source]

Bases: NamedTuple

matvec_max: int | None: Alias for field number 0

abstol: float: Alias for field number 1

reltol: float: Alias for field number 2

precon: Tensor | Callable[[Tensor], Tensor] | None: Alias for field number 3

logger: Logger: Alias for field number 4

torchsparsegradutils.utils.bicgstab.bicgstab(matmul_closure: ~torch.Tensor | ~typing.Callable[[~torch.Tensor], ~torch.Tensor], rhs: ~torch.Tensor, initial_guess: ~torch.Tensor | None = None, settings: ~torchsparsegradutils.utils.bicgstab.BICGSTABSettings = (None, 1e-08, 1e-06, None, <Logger bicgstab (WARNING)>)) → Tensor[source]

Solve linear systems with the BiConjugate Gradient Stabilized (BiCGSTAB) method.

Solves nonsymmetric, nonsingular systems \(A x = b\). Accepts either a matrix-like tensor (using .matmul) or a callable for the matrix–vector product, and optionally a (left) preconditioner, also as tensor or callable, approximating \(M^{-1}\).

Parameters:

matmul_closure ({torch.Tensor, callable(x) -> Ax}) – Matrix–vector multiplication operator. If a tensor is provided, its .matmul is used.
rhs (torch.Tensor, shape (n,) or (n, k)) – Right-hand side vector(s). For multiple RHS, each column is solved independently (no block BiCGSTAB).
initial_guess (torch.Tensor, optional, shape like rhs) – Initial guess. If None, zero initialization is used.
settings (BICGSTABSettings, optional) – Convergence tolerances, maximum matvecs, optional preconditioner and logger.

Returns:

Solution(s) x with the same shape as ``rhs``.

Return type:

torch.Tensor

Raises:

RuntimeError – If matmul_closure is neither tensor nor callable, or if the precon is neither tensor nor callable.

Notes

Per iteration (unpreconditioned) BiCGSTAB [1a] uses ~2 matvecs, several dot products, and vector updates. The algorithm can experience breakdown when certain inner products or denominators vanish (e.g., \(\langle r_0, v \rangle = 0\) or \(\langle t, t \rangle = 0\)). This implementation follows a standard variant [2a] and solves multiple RHS by looping over columns (no shared Krylov subspace).

References

[1a]

Van der Vorst, H. A. (1992). Bi-CGSTAB: A fast and smoothly converging variant of Bi-CG for the solution of nonsymmetric linear systems. SIAM J. Sci. Stat. Comput., 13(2), 631–644.

[2a]

Kelley, C. T. (1995). Iterative Methods for Linear and Nonlinear Equations. SIAM.

Examples

>>> import torch
>>> from torchsparsegradutils.utils import bicgstab
>>> A = torch.tensor([[3.0, 1.0], [2.0, 4.0]])
>>> b = torch.tensor([1.0, 2.0])
>>> x = bicgstab(A.matmul, b)
>>> x.shape
torch.Size([2])

Multiple right-hand sides:

>>> B = torch.randn(2, 3)
>>> X = bicgstab(A.matmul, B)
>>> X.shape
torch.Size([2, 3])

With custom settings:

>>> from torchsparsegradutils.utils.bicgstab import BICGSTABSettings
>>> settings = BICGSTABSettings(abstol=1e-10, reltol=1e-8, matvec_max=1000)
>>> x = bicgstab(A.matmul, b, settings=settings)

With preconditioning:

>>> # Diagonal preconditioner
>>> # Extract and regularize diagonal
>>> diagA = torch.diag(A)
>>> eps = 1e-12
>>> safe_diag = torch.where(diagA.abs() < eps, torch.full_like(diagA, eps), diagA)
>>> inv_diag = 1.0 / safe_diag
>>> # Supply as an operator (apply M^{-1} r = inv_diag * r elementwise)
>>> settings_precond = BICGSTABSettings(
...     precon=lambda r: inv_diag * r  # r has same shape as b
... )
>>> x = bicgstab(A.matmul, b, settings=settings_precond)

class torchsparsegradutils.utils.linear_cg.LinearCGSettings(max_cg_iterations, max_lanczos_quadrature_iterations, cg_tolerance, terminate_cg_by_size, verbose_linalg)[source]

Bases: NamedTuple

max_cg_iterations: int: Alias for field number 0

max_lanczos_quadrature_iterations: int: Alias for field number 1

cg_tolerance: float: Alias for field number 2

terminate_cg_by_size: bool: Alias for field number 3

verbose_linalg: bool: Alias for field number 4

torchsparsegradutils.utils.linear_cg.linear_cg(matmul_closure: Tensor | Callable[[Tensor], Tensor], rhs: Tensor, n_tridiag: int = 0, tolerance: float | None = None, eps: float = 1e-10, stop_updating_after: float = 1e-10, max_iter: int | None = None, max_tridiag_iter: int | None = None, initial_guess: Tensor | None = None, preconditioner: Callable[[Tensor], Tensor] | None = None, settings: LinearCGSettings = (1000, 20, 1, False, False)) → Tensor | tuple[Tensor, Tensor][source]

Solve symmetric positive definite linear systems using conjugate gradient (CG).

Implements linear CG for systems \(A x = b\) with symmetric positive definite operator \(A\). Supports single/multiple RHS and optional stochastic Lanczos tridiagonalization for eigenvalue/log-determinant estimation.

Parameters:

matmul_closure ({torch.Tensor, callable(x) -> A x}) – Matrix–vector multiply. If a tensor is provided, its .matmul is used. The callable must accept inputs shaped like rhs and return A @ rhs.
rhs (torch.Tensor, shape (..., n) or (..., n, k)) – Right-hand side(s). Leading batch dims are supported.
n_tridiag (int, optional) – Number of Lanczos tridiagonalizations (probe vectors). If > 0, tridiagonal matrices are returned in addition to the solution. Default: 0.
tolerance (float, optional) – Average residual-norm stopping criterion. If None, uses settings.cg_tolerance.
eps (float, optional) – Small constant to avoid division by zero. Default: 1e-10.
stop_updating_after (float, optional) – Per-vector early-stop threshold for residual norms. Default: 1e-10.
max_iter (int, optional) – Maximum CG iterations. If None, uses settings.max_cg_iterations.
max_tridiag_iter (int, optional) – Maximum Lanczos size. If None, uses settings.max_lanczos_quadrature_iterations.
initial_guess (torch.Tensor, optional, shape like rhs) – Initial guess. If None, zeros are used.
preconditioner (callable, optional) – Preconditioner with signature preconditioner(x) -> M^{-1} x. If None, no preconditioning is used.
settings (LinearCGSettings, optional) – Configuration for iteration caps, tolerances, and logging verbosity.

Returns:

If n_tridiag == 0: solution x with the same shape as rhs.
If n_tridiag > 0: (x, T) where T has shape (n_tridiag, *rhs.shape[:-2], r, r) with r = last_tridiag_iter + 1 and r <= min(max_tridiag_iter, n). Without batch dims this is (n_tridiag, r, r).

Return type:

torch.Tensor or (torch.Tensor, torch.Tensor)

Raises:

RuntimeError – If max_tridiag_iter > max_iter.
RuntimeError – If matmul_closure is neither a tensor nor a callable.

Notes

CG converges in at most n iterations for SPD matrices, but typically much faster if eigenvalues are clustered. Preconditioning (e.g. diagonal or incomplete Cholesky) can significantly accelerate convergence. When n_tridiag > 0, Lanczos tridiagonalization is accumulated alongside CG for spectral / log-determinant estimates.

This implementation is based on MIT-licensed code from the linear_operator library [1e].

Examples

Basic CG solve:

>>> A = torch.tensor([[4.0, -1.0], [-1.0, 4.0]])
>>> b = torch.tensor([1.0, 2.0])
>>> x = linear_cg(A.matmul, b)
>>> x.shape
torch.Size([2])

Multiple RHS:

>>> B = torch.randn(2, 5)  # 5 RHS
>>> X = linear_cg(A.matmul, B, max_iter=100, tolerance=1e-8)
>>> X.shape
torch.Size([2, 5])

With preconditioning:

>>> M_inv = torch.diag(1.0 / torch.diag(A))
>>> x = linear_cg(A.matmul, b, preconditioner=lambda v: M_inv @ v)

With Lanczos tridiagonalization:

>>> x, T = linear_cg(A.matmul, b, n_tridiag=1)
>>> T.shape  # (n_tridiag, r, r) with r <= n
torch.Size([1, 2, 2])

Sparse operator via closure:

>>> indices = torch.tensor([[0, 0, 1, 1, 2], [0, 1, 0, 1, 2]])
>>> values = torch.tensor([4.0, -1.0, -1.0, 4.0, 2.0])
>>> A_sp = torch.sparse_coo_tensor(indices, values, (3, 3))
>>> x = linear_cg(lambda v: A_sp @ v, torch.randn(3))

References

[1e]

linear_operator library. https://github.com/cornellius-gp/linear_operator

Code adapted from https://github.com/rfeinman/pytorch-minimize/blob/master/torchmin/lstsq/lsmr.py Code modified from scipy.sparse.linalg.lsmr

torchsparsegradutils.utils.lsmr.lsmr(A: Tensor | Callable[[Tensor], Tensor], b: Tensor, Armat: Tensor | Callable[[Tensor], Tensor] | None = None, n: int | None = None, damp: float = 0.0, atol: float = 1e-06, btol: float = 1e-06, conlim: float = 100000000.0, maxiter: int | None = None, x0: Tensor | None = None, check_nonzero: bool = True) → Tuple[Tensor, int][source]

Least Squares Minimal Residual (LSMR) solver.

Iterative solver for \(A x = b\) and least-squares problems \(\min_x \|A x - b\|_2\). Works with large, sparse, or rectangular \(A\) and is often more stable than LSQR on ill-conditioned problems.

Parameters:

A ({torch.Tensor, callable(x) -> A @ x}) – System matrix or matvec closure. If a tensor is given, it may be dense or sparse and .matmul is used.
b (torch.Tensor, shape (m,)) – Right-hand side vector. Must be on the same device/dtype as A.
Armat ({torch.Tensor, callable(x) -> A^T @ x}, optional) – Transpose matvec or matrix. If A is a tensor and Armat is None, uses A.adjoint().matmul. If A is callable, Armat is required.
n (int, optional) – Number of columns of A. Required if A is callable; inferred from A.shape[1] if A is a tensor.
damp (float, optional) – Tikhonov damping parameter (ridge). Solves \(\min_x \|(A; \text{damp} I) x - (b; 0)\|_2\). Default: 0.0.
atol (float, optional) – Absolute convergence tolerance. Default: 1e-6.
btol (float, optional) – Relative residual tolerance. Default: 1e-6.
conlim (float, optional) – Condition estimate limit; stops if estimate exceeds this value. Default: 1e8.
maxiter (int, optional) – Maximum iterations. If None, uses min(m, n).
x0 (torch.Tensor, optional, shape (n,)) – Initial guess. If None, zeros are used.
check_nonzero (bool, optional) – Skip the rare beta == 0 synchronization check for performance when set to False (use with caution). Default: True.

Returns:

x (torch.Tensor, shape (n,)) – Approximate solution that minimizes \(\|A x - b\|_2\) (with damped variant when damp > 0).
iterations (int) – Number of iterations executed.

Raises:

RuntimeError – If A is neither a tensor nor a callable.
RuntimeError – If A is callable and n is not provided.
RuntimeError – If Armat is missing or is neither a tensor nor a callable.

Notes

Uses Golub–Kahan bidiagonalization [1f] with specialized QR steps. For overdetermined systems (m > n), returns the least-squares solution. For underdetermined systems (m < n) with damp = 0, returns the minimum-norm least-squares solution.

Convergence checks (roughly):

Consistent: \(\|r\|_2 \le \text{atol} \, \|A\| \, \|x\| + \text{btol} \, \|b\|\)

Inconsistent: \(\|A^\top r\|_2 \le \text{atol} \, \|A\| \, \|r\|\)

References

[1f]

Fong, D. C., & Saunders, M. (2011). LSMR: An iterative algorithm for sparse least-squares problems. SIAM Journal on Scientific Computing, 33(5), 2950-2971.

Examples

Basic least squares problem:

>>> import torch
>>> from torchsparsegradutils.utils import lsmr
>>> # Over-determined system (3x2)
>>> A = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
>>> b = torch.tensor([1.0, 2.0, 3.0])
>>> x, iterations = lsmr(A, b)
>>> x.shape
torch.Size([2])

Sparse matrix least squares:

>>> # Create sparse matrix
>>> indices = torch.tensor([[0, 1, 2, 1, 2], [0, 0, 0, 1, 1]])
>>> values = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
>>> A_sparse = torch.sparse_coo_tensor(indices, values, (3, 2))
>>> x, it = lsmr(A_sparse, b)

With damping for regularization (Tikhonov / ridge):

>>> # Regularized least squares
>>> x_reg, it = lsmr(A, b, damp=0.1)

Callable matrix interface:

>>> x, it = lsmr(lambda v: A @ v, b, Armat=lambda v: A.T @ v, n=2)

Under-determined system (minimum norm solution):

>>> A_under = torch.randn(2, 4)  # 2x4 system
>>> b_under = torch.randn(2)
>>> x_min_norm, it = lsmr(A_under, b_under)
>>> x_min_norm.shape
torch.Size([4])

Custom tolerances and limits:

>>> x, it = lsmr(A, b, atol=1e-10, btol=1e-10, conlim=1e12, maxiter=1000)

class torchsparsegradutils.utils.minres.MINRESSettings(max_cg_iterations, minres_tolerance, verbose_linalg)[source]

Bases: NamedTuple

max_cg_iterations: int: Alias for field number 0

minres_tolerance: float: Alias for field number 1

verbose_linalg: bool: Alias for field number 2

torchsparsegradutils.utils.minres.minres(matmul_closure: Tensor | Callable[[Tensor], Tensor], rhs: Tensor, eps: float = 1e-25, shifts: Tensor | None = None, value: float | None = None, max_iter: int | None = None, preconditioner: Callable[[Tensor], Tensor] | None = None, settings: MINRESSettings = (1000, 0.0001, False)) → Tensor[source]

Minimum Residual (MINRES) solver for symmetric (Hermitian) linear systems.

Solves linear systems A x = b where A is symmetric (Hermitian) and may be indefinite. Supports single/multiple right-hand sides and (optionally) multiple shift values to solve (A + \sigma I) x = b in one run. Gradually minimizes the residual norm ||A x - b||_2 via the Lanczos process.

Parameters:

matmul_closure ({torch.Tensor, callable(x) -> A @ x}) – Matrix–vector multiplication operator. If a tensor is provided, its .matmul is used. The operator should represent a symmetric/Hermitian matrix for MINRES to behave as intended.
rhs (torch.Tensor, shape (..., n) or (..., n, k)) – Right-hand side vector(s). Leading batch dimensions are supported; for multi-RHS, the last two dims are (n, k).
eps (float, optional) – Small constant to prevent division by zero/numerical issues. Default: 1e-25.
shifts (torch.Tensor or scalar, optional) – Shift(s) \sigma for solving (A + \sigma I) x = b. If None or a scalar, a single system is solved. If a tensor with s elements, the solver computes s shifted systems and stacks their solutions along a new leading dimension.
value (float, optional) – Scalar multiplier \alpha applied to the operator (solves (\alpha A) x = b) when provided. Default: None (no scaling).
max_iter (int, optional) – Maximum iterations. If None, uses settings.max_cg_iterations. Internally capped at n + 1 where n is the problem size.
preconditioner (callable, optional) – Left preconditioner with signature preconditioner(x) -> M^{-1} x. If None, no preconditioning is used.
settings (MINRESSettings, optional) – Configuration object controlling iteration caps and tolerances (e.g., minres_tolerance for the relative update criterion).

Returns:

If shifts is None or a scalar: solution with the same shape as rhs (i.e., (..., n) or (..., n, k)). If shifts has length s: a stacked tensor of shape (s, *rhs.shape) containing solutions for each shift.

Return type:

torch.Tensor

Raises:

RuntimeError – If matmul_closure is neither a tensor nor a callable.

Notes

MINRES [1g] is appropriate for symmetric/Hermitian indefinite systems; it minimizes the Euclidean residual norm rather than the A-norm (as in CG).
For symmetric positive definite systems, Conjugate Gradient (CG) typically converges faster; prefer CG unless indefiniteness/robustness suggests MINRES.
When multiple shifts are provided, the solver reuses Lanczos information and returns one solution per shift value.
All inputs should share device and dtype; the implementation normalizes rhs internally and rescales the final solution(s).

See also

linear_cg: Conjugate Gradient for SPD systems.
bicgstab: BiCGSTAB for general non-symmetric systems.

References

[1g]

Paige, C. C., & Saunders, M. A. (1975). Solution of sparse indefinite systems of linear equations. SIAM Journal on Numerical Analysis, 12(4), 617–629.

Examples

Basic solve (indefinite, symmetric):

>>> A = torch.tensor([[2.0, 1.0], [1.0, -1.0]])
>>> b = torch.tensor([1.0, 2.0])
>>> x = minres(A.matmul, b)
>>> x.shape
torch.Size([2])

Multiple right-hand sides:

>>> B = torch.randn(2, 3)
>>> X = minres(A.matmul, B)
>>> X.shape
torch.Size([2, 3])

Shifted system (regularization):

>>> x_shifted = minres(A.matmul, b, shifts=torch.tensor(0.1))

Sparse operator via closure:

>>> idx = torch.tensor([[0, 0, 1, 1], [0, 1, 0, 1]])
>>> val = torch.tensor([2.0, 1.0, 1.0, -1.0])
>>> A_sp = torch.sparse_coo_tensor(idx, val, (2, 2))
>>> x = minres(lambda v: A_sp @ v, b)

With a simple diagonal preconditioner:

>>> M_diag = torch.abs(torch.diag(A)) + 0.1
>>> precond = lambda x: x / M_diag.unsqueeze(-1)
>>> x = minres(A.matmul, b, preconditioner=precond)

Custom iteration cap/tolerance:

>>> settings = MINRESSettings(max_cg_iterations=200, minres_tolerance=1e-5)
>>> x = minres(A.matmul, b, settings=settings)

Sparse Utilities

torchsparsegradutils.utils.utils.stack_csr(tensors: List[Tensor], dim: int = 0) → Tensor[source]

Stack CSR sparse tensors along a new dimension.

This function is analogous to torch.stack(), but specifically designed for CSR (Compressed Sparse Row) tensors. Unlike COO tensors, CSR tensors are not currently supported by torch.stack(), hence this helper provides the missing functionality.

Parameters:

tensors (list of torch.Tensor) – List of 2D CSR sparse tensors to be stacked. All tensors must have the same shape and layout.
dim (int, default=0) – Dimension along which to stack the tensors.

Returns:

A CSR sparse tensor with an additional dimension of size len(tensors) inserted at position dim.

Return type:

torch.Tensor

Raises:

TypeError – If tensors is not a list or tuple.
ValueError – If tensors is empty, contain tensors of different shapes, are not in CSR format, or are not 2D.

Notes

torch.stack() supports COO sparse tensors but not CSR. This function fills that gap by implementing stacking logic for CSR tensors.

Examples

Stack multiple 2D CSR tensors:

>>> import torch
>>> from torchsparsegradutils.utils import stack_csr
>>> crow = torch.tensor([0, 1, 2])
>>> col = torch.tensor([0, 1])
>>> A = torch.sparse_csr_tensor(crow, col, torch.tensor([1.0, 2.0]), (2, 2))
>>> B = torch.sparse_csr_tensor(crow, col, torch.tensor([3.0, 4.0]), (2, 2))
>>> stacked = stack_csr([A, B])
>>> stacked.shape
torch.Size([2, 2, 2])

The new dimension is a CSR batch dimension:

>>> stacked.crow_indices().shape
torch.Size([2, 3])

torchsparsegradutils.utils.utils.convert_coo_to_csr_indices_values(coo_indices, num_rows, values=None)[source]

Convert COO indices to CSR format with optional value permutation.

Converts COO (Coordinate List Format) row and column indices to CSR (Compressed Sparse Row) format. Supports both batched and unbatched indices. The function sorts the COO indices lexicographically and compresses row indices to CSR crow format.

Parameters:

coo_indices (torch.Tensor) – COO indices tensor with shape (2, nnz) for unbatched or (3, nnz) for batched format. Rows are [row_idx, col_idx] or [batch_idx, row_idx, col_idx].
num_rows (int) – Number of rows in the matrix.
values (torch.Tensor, optional) – Values tensor corresponding to COO indices. If provided, values are reordered according to the index sorting permutation.

Returns:

crow_indices (torch.Tensor) – CSR crow indices. For unbatched: shape (num_rows + 1,). For batched: shape (num_batches, num_rows + 1).
col_indices (torch.Tensor) – CSR column indices. Same shape as original column indices but reordered according to sorting.
values_or_permutation (torch.Tensor) – If values provided: reordered values tensor. If values is None: permutation indices from sorting.

Raises:

ValueError – If indices tensor has wrong number of dimensions, row indices exceed num_rows, or values shape doesn’t match indices.

Examples

Unbatched COO to CSR conversion:

>>> import torch
>>> from torchsparsegradutils.utils import convert_coo_to_csr_indices_values
>>> # COO indices for 3x3 matrix
>>> coo_indices = torch.tensor([[0, 1, 2], [1, 0, 2]])
>>> values = torch.tensor([1.0, 2.0, 3.0])
>>> crow, col, vals = convert_coo_to_csr_indices_values(
...     coo_indices, num_rows=3, values=values)
>>> crow
tensor([0, 1, 2, 3])

Batched conversion:

>>> # Batched COO indices: 2 batches, each with 2 elements
>>> batch_coo = torch.tensor([[0, 0, 1, 1], [0, 1, 0, 1], [1, 0, 1, 0]])
>>> batch_values = torch.tensor([1.0, 2.0, 3.0, 4.0])
>>> crow, col, vals = convert_coo_to_csr_indices_values(
...     batch_coo, num_rows=2, values=batch_values)

Without values (get permutation):

>>> crow, col, perm = convert_coo_to_csr_indices_values(
...     coo_indices, num_rows=3)

torchsparsegradutils.utils.utils.convert_coo_to_csr(sparse_coo_tensor)[source]

Convert COO sparse tensor to CSR format.

Converts a COO (Coordinate List Format) sparse tensor to CSR (Compressed Sparse Row) format. Handles both unbatched and batched tensors with optional leading batch dimension.

Parameters:: sparse_coo_tensor (torch.Tensor) – COO sparse tensor to convert. Must have layout torch.sparse_coo. Can be 2D (m, n) or 3D (b, m, n) with single batch dimension.
Returns:: CSR sparse tensor with same shape and values as input.
Return type:: torch.Tensor
Raises:: ValueError – If input tensor layout is not torch.sparse_coo.

Notes

The function automatically coalesces the COO tensor if not already coalesced, ensuring proper handling of duplicate indices.

Examples

Convert 2D COO to CSR:

>>> import torch
>>> from torchsparsegradutils.utils import convert_coo_to_csr
>>> # Create COO tensor
>>> indices = torch.tensor([[0, 1, 1], [1, 0, 2]])
>>> values = torch.tensor([1.0, 2.0, 3.0])
>>> coo_tensor = torch.sparse_coo_tensor(indices, values, (2, 3))
>>> csr_tensor = convert_coo_to_csr(coo_tensor)
>>> csr_tensor.layout
torch.sparse_csr

Convert batched COO to CSR:

>>> # Batched COO tensor: 2 batches, each with 2 elements
>>> batch_indices = torch.tensor([[0, 0, 1, 1], [0, 1, 0, 1], [1, 0, 1, 0]])
>>> batch_values = torch.tensor([1.0, 2.0, 3.0, 4.0])
>>> batched_coo = torch.sparse_coo_tensor(batch_indices, batch_values, (2, 2, 2))
>>> batched_csr = convert_coo_to_csr(batched_coo)
>>> batched_csr.shape
torch.Size([2, 2, 2])

torchsparsegradutils.utils.utils.sparse_block_diag(*sparse_tensors: Tensor) → Tensor[source]

Construct a block-diagonal sparse matrix from COO/CSR inputs.

Builds a block-diagonal sparse tensor from a sequence of 2D sparse tensors that are all in the same layout (either COO or CSR). The result has blocks placed on the diagonal and zeros elsewhere, analogous to torch.block_diag() but for sparse inputs.

Parameters:

*sparse_tensors (torch.Tensor) – Variable number of 2D sparse tensors (all COO or all CSR). Each tensor must have exactly 2 sparse dimensions and 0 dense dimensions.

Returns:

A sparse tensor in the same layout as the inputs with shape (sum_i n_i, sum_i m_i), where each block i has shape (n_i, m_i).

Return type:

torch.Tensor

Raises:

TypeError – If any input is not a torch.Tensor.
ValueError – If no tensors are provided; if layouts are mixed; or if any input does not have exactly 2 sparse dims and 0 dense dims.

Notes

The resulting block structure is

[A₁  0   0  ... 0 ]
[0   A₂  0  ... 0 ]
[0   0   A₃ ... 0 ]
[⋮   ⋮   ⋮  ⋱  ⋮ ]
[0   0   0  ... Aₙ]

Offsets for each block are computed using cumulative row/column sizes of all preceding blocks (not simply i * size), so inputs may have different shapes.

Examples

COO inputs:

>>> import torch
>>> from torchsparsegradutils.utils import sparse_block_diag
>>> A = torch.sparse_coo_tensor(torch.tensor([[0, 1], [0, 1]]), torch.tensor([1., 2.]), size=(2, 2))
>>> B = torch.sparse_coo_tensor(torch.tensor([[0], [0]]), torch.tensor([3.]), size=(1, 1))
>>> C = sparse_block_diag(A, B)
>>> C.shape
torch.Size([3, 3])
>>> C.layout
torch.sparse_coo

CSR inputs:

>>> A_csr = A.to_sparse_csr()
>>> B_csr = B.to_sparse_csr()
>>> D = sparse_block_diag(A_csr, B_csr)
>>> D.layout
torch.sparse_csr

See also

torch.block_diag: Dense block-diagonal construction for dense inputs.
stack_csr: Stack CSR matrices along a new batch dimension.

torchsparsegradutils.utils.utils.sparse_block_diag_split(sparse_block_diag_tensor: Tensor, *shapes: Tuple[int, int]) → tuple[Tensor, ...][source]

Split a block-diagonal sparse matrix back into its component blocks.

Given a block-diagonal sparse tensor produced by sparse_block_diag(), return the original 2D sparse tensors (in the same layout) according to the provided shapes. Supports COO and CSR layouts.

Parameters:

sparse_block_diag_tensor (torch.Tensor) – Input block-diagonal sparse tensor (COO or CSR). Must be 2D and have exactly two sparse dimensions and zero dense dimensions.
*shapes (tuple of int) – Sequence of shapes (rows_i, cols_i) for each block in the order they appear along the diagonal. The sums of rows and cols must match the input tensor’s height and width, respectively.

Returns:

The recovered sparse blocks, each a 2D sparse tensor in the same layout as sparse_block_diag_tensor.

Return type:

tuple of torch.Tensor

Raises:

ValueError – If the input layout is not COO or CSR; if any provided shape is not 2D; or if the sum of provided shapes does not match the input size.
TypeError – If sparse_block_diag_tensor is not a tensor.

Notes

For COO inputs, this function assumes the tensor is coalesced. If it is not, it will be coalesced internally to avoid duplicate coordinates.
This is the inverse operation of sparse_block_diag() when given the correct shapes (order and sizes) of the original blocks.

See also

sparse_block_diag: Construct a block-diagonal sparse matrix from 2D sparse blocks.

torchsparsegradutils.utils.utils.sparse_eye(size: Tuple[int, ...], *, layout: layout = torch.sparse_coo, values_dtype: dtype = torch.float64, indices_dtype: dtype = torch.int64, device: device = device(type='cpu'), requires_grad: bool = False) → Tensor[source]

Create a sparse identity matrix.

Constructs an identity matrix in sparse format (COO or CSR). Supports both unbatched and batched square matrices.

Parameters:

size (tuple of int) – Shape of the identity matrix. Must be either (n, n) for unbatched or (batch_size, n, n) for batched. Rows and columns must be equal.
layout (torch.layout, default=torch.sparse_coo) – Sparse tensor layout. Must be either torch.sparse_coo or torch.sparse_csr.
values_dtype (torch.dtype, default=torch.float64) – Data type of the values. Only torch.float32 and torch.float64 are supported.
indices_dtype (torch.dtype, default=torch.int64) – Data type of the indices. Must be torch.int32 or torch.int64.
device (torch.device, default=torch.device("cpu")) – Device on which to create the tensor.
requires_grad (bool, default=False) – Whether autograd should record operations on the returned tensor.

Returns:

Sparse identity matrix of shape (n, n) or batched identity matrix of shape (batch_size, n, n), in the requested sparse layout.

Return type:

torch.Tensor

Raises:

ValueError – If size is not 2D or 3D, matrix is not square, or if dtypes/layout are unsupported.

Notes

For batched inputs, each batch element is an independent identity matrix.

Examples

Unbatched identity (COO):

>>> from torchsparsegradutils.utils import sparse_eye
>>> I = sparse_eye((3, 3), layout=torch.sparse_coo)
>>> I.to_dense()  
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]]...)

Batched identity (CSR):

>>> I_batched = sparse_eye((2, 4, 4), layout=torch.sparse_csr)
>>> I_batched.shape
torch.Size([2, 4, 4])

Random Sparse Tensor Generation

Utility functions for generating random sparse matrices.

Notes

Sparse COO tensors have indices tensor of size (ndim, nse) with index dtype torch.int64.
Sparse CSR tensors store structure via crow_indices and col_indices whose dtype may be
torch.int64 (default) or torch.int32. If you want MKL-enabled matrix operations, prefer torch.int32 (PyTorch is typically linked with MKL LP64 which uses 32-bit integer indexing).
Batched sparse CSR tensors currently require each batch to have the same number of specified elements;
this constraint enables efficient storage of batched CSR indices.

torchsparsegradutils.utils.random_sparse.rand_sparse(size: Tuple[int, int] | Tuple[int, int, int], nnz: int, layout: layout = torch.sparse_coo, *, indices_dtype: dtype = torch.int64, values_dtype: dtype = torch.float32, device: device = device(type='cpu'), well_conditioned: bool = False, min_diag_value: float = 1.0) → Tensor[source]

Generate a random sparse matrix.

A convenience wrapper around generate_random_sparse_coo_matrix() and generate_random_sparse_csr_matrix(), dispatching based on the requested layout.

Parameters:

size (tuple of int) – Shape of the matrix, either (n_r, n_c) or (b, n_r, n_c).
nnz (int) – Number of nonzeros per batch item.
layout (torch.layout, default=torch.sparse_coo) – Sparse format. Must be torch.sparse_coo or torch.sparse_csr.
indices_dtype (torch.dtype, default=torch.int64) – Index dtype (torch.int64 or torch.int32).
values_dtype (torch.dtype, default=torch.float32) – Data type of the nonzero values.
device (torch.device, default=torch.device("cpu")) – Device for tensor allocation.
well_conditioned (bool, default=False) – If True and square, diagonal values are boosted for stability. See generate_random_sparse_coo_matrix().
min_diag_value (float, default=1.0) – Minimum diagonal value when well_conditioned=True.

Returns:

A sparse COO or CSR tensor.

Return type:

torch.Tensor

Raises:

ValueError – If layout is not supported.

See also

generate_random_sparse_coo_matrix: Generate a random sparse COO matrix.
generate_random_sparse_csr_matrix: Generate a random sparse CSR matrix.

Examples

>>> A = rand_sparse((100, 100), 500)
>>> A.layout
torch.sparse_coo
>>> B = rand_sparse((50, 50), 200, layout=torch.sparse_csr, well_conditioned=True)
>>> B.layout
torch.sparse_csr

torchsparsegradutils.utils.random_sparse.rand_sparse_tri(size: Tuple[int, int] | Tuple[int, int, int], nnz: int, layout: layout = torch.sparse_coo, *, upper: bool = True, strict: bool = False, indices_dtype: dtype = torch.int64, values_dtype: dtype = torch.float32, device: device = device(type='cpu'), value_range: Tuple[float, float] = (0, 1), well_conditioned: bool = False, min_diag_value: float = 1.0) → Tensor[source]

Generate a random sparse triangular matrix.

A convenience wrapper around the triangular sparse matrix generators: - generate_random_sparse_triangular_coo_matrix() - generate_random_sparse_triangular_csr_matrix() - generate_random_sparse_strictly_triangular_coo_matrix() - generate_random_sparse_strictly_triangular_csr_matrix()

Parameters:

size (tuple of int) – Shape of the square matrix, (n, n) or (b, n, n).
nnz (int) – Number of nonzeros per batch item. - If strict=True: nnz <= n*(n-1)/2. - If strict=False: n <= nnz <= n*(n+1)/2 (includes diagonal).
layout (torch.layout, default=torch.sparse_coo) – Sparse format. Must be torch.sparse_coo or torch.sparse_csr.
upper (bool, default=True) – If True, generate upper-triangular. If False, lower-triangular.
strict (bool, default=False) – If True, exclude diagonal. If False, include diagonal.
indices_dtype (torch.dtype, default=torch.int64) – Index dtype.
values_dtype (torch.dtype, default=torch.float32) – Data type of nonzero values.
device (torch.device, default=torch.device("cpu")) – Device for tensor allocation.
value_range (tuple of float, default=(0, 1)) – Range for random values.
well_conditioned (bool, default=False) – If True, diagonal values are boosted. Only used when strict=False.
min_diag_value (float, default=1.0) – Minimum diagonal value when well_conditioned=True.

Returns:

A sparse triangular COO or CSR tensor.

Return type:

torch.Tensor

Raises:

ValueError – If layout is not supported.
ValueError – If nnz does not satisfy the constraints for strict/non-strict.

Examples

>>> A = rand_sparse_tri((100, 100), 500)
>>> A.layout
torch.sparse_coo

See also

generate_random_sparse_triangular_coo_matrix: Non-strict triangular COO generator.
generate_random_sparse_triangular_csr_matrix: Non-strict triangular CSR generator.
generate_random_sparse_strictly_triangular_coo_matrix: Strict triangular COO generator.
generate_random_sparse_strictly_triangular_csr_matrix: Strict triangular CSR generator.

torchsparsegradutils.utils.random_sparse.make_spd_sparse(n: int, layout: layout, value_dtype: dtype, index_dtype: dtype, device: device, sparsity_ratio: float = 0.5, nz: int | None = None) → Tuple[Tensor, Tensor][source]

Generate a random sparse symmetric positive definite (SPD) matrix.

Constructs a dense SPD matrix A = M Mᵀ + n I from a random Gaussian matrix M, then applies structured sparsification by zeroing out symmetric pairs of off-diagonal entries. Converts the result to the requested sparse layout.

Parameters:

n (int) – Dimension of the matrix (produces an n × n SPD matrix).
layout (torch.layout) – Sparse tensor layout to return. Must be torch.sparse_coo or torch.sparse_csr.
value_dtype (torch.dtype) – Data type for matrix values.
index_dtype (torch.dtype) – Data type for sparse indices. Typically torch.int32 or torch.int64. Note: PyTorch may coerce COO indices to int64 on coalesce, even if int32 is requested. CSR tensors preserve the chosen dtype.
device (torch.device) – Device on which to allocate the tensors (e.g., torch.device("cuda")).
sparsity_ratio (float, optional) – Fraction of off-diagonal upper-triangular elements to zero out. Each selection removes both (i, j) and (j, i) to maintain symmetry. Ignored if nz is provided. Default is 0.5.
nz (int, optional) – Exact number of symmetric pairs of off-diagonal elements to zero out. Each pair corresponds to (i, j) and (j, i) for i ≠ j. If None, sparsity_ratio is used. Default is None.

Returns:

A_sparse (torch.Tensor) – Sparse SPD matrix in the requested layout (torch.sparse_coo or torch.sparse_csr) with shape (n, n).
A_dense (torch.Tensor) – Dense SPD matrix with shape (n, n) before conversion to sparse.

Raises:

ValueError – If layout is not torch.sparse_coo or torch.sparse_csr.

Notes

The SPD property is guaranteed by construction as A = M Mᵀ + n I, regardless of sparsification.
Sparsification is symmetric: whenever entry (i, j) is zeroed, (j, i) is also zeroed, preserving symmetry.
Only unbatched (2D) matrices are supported. For batched sparse SPD matrices, extend this function accordingly.
For COO tensors, PyTorch may coerce indices to int64 during coalesce. For CSR tensors, the requested index_dtype is preserved.

Examples

Generate a sparse SPD matrix in COO format:

>>> A_sp, A_dn = make_spd_sparse(
...     n=5,
...     layout=torch.sparse_coo,
...     value_dtype=torch.float32,
...     index_dtype=torch.int64,
...     device=torch.device("cpu"),
...     sparsity_ratio=0.6,
... )
>>> A_sp.shape
torch.Size([5, 5])
>>> A_dn.shape
torch.Size([5, 5])

Generate a sparse SPD matrix in CSR format with exactly 4 zeroed pairs:

>>> A_sp, A_dn = make_spd_sparse(
...     n=6,
...     layout=torch.sparse_csr,
...     value_dtype=torch.float64,
...     index_dtype=torch.int32,
...     device=torch.device("cpu"),
...     nz=4,
... )

torchsparsegradutils.utils.random_sparse.generate_random_sparse_coo_matrix(size: Tuple[int, int] | Tuple[int, int, int], nnz: int, *, indices_dtype: dtype = torch.int64, values_dtype: dtype = torch.float32, device: device = device(type='cpu'), well_conditioned: bool = False, min_diag_value: float = 1.0) → Tensor[source]

Generate a random sparse COO matrix.

Creates an unbatched (n_r, n_c) or batched (b, n_r, n_c) COO tensor with exactly nnz nonzeros per batch item, sampled uniformly at random across all entries. Optionally boosts diagonal entries for square matrices when well_conditioned=True.

Parameters:

size (tuple of int) – Shape of the output matrix. Must be (n_r, n_c) or (b, n_r, n_c).
nnz (int) – Number of nonzero entries per batch item. Must satisfy 0 <= nnz <= n_r * n_c.
indices_dtype (torch.dtype, default=torch.int64) – Index dtype for the COO indices (torch.int64 or torch.int32).
values_dtype (torch.dtype, default=torch.float32) – Data type of the nonzero values.
device (torch.device, default=torch.device("cpu")) – Device to allocate the tensors on.
well_conditioned (bool, default=False) – If True and the matrix is square, diagonal entries (if present among the sampled nonzeros) are boosted to be at least min_diag_value.
min_diag_value (float, default=1.0) – Minimum diagonal value when well_conditioned=True.

Returns:

A sparse COO tensor of shape size with nnz nonzeros per batch (or exactly nnz for the unbatched case).

Return type:

torch.Tensor

Raises:

ValueError – If size has fewer than 2 dims, more than 3 dims, or nnz is out of range.
ValueError – If indices_dtype is not torch.int64 or torch.int32.

Notes

Nonzeros are placed uniformly at random across all entries.

See also

rand_sparse: Convenience dispatcher that selects COO/CSR generator based on layout.
generate_random_sparse_triangular_coo_matrix: Generate triangular COO matrices.
generate_random_sparse_strictly_triangular_coo_matrix: Generate strictly triangular COO matrices.

Examples

Unbatched random COO (200 nonzeros):

>>> A = generate_random_sparse_coo_matrix((100, 50), 200)

Batched random COO (b=3, each with 500 nonzeros):

>>> A = generate_random_sparse_coo_matrix((3, 100, 100), 500)

Well-conditioned square matrix:

>>> A = generate_random_sparse_coo_matrix((100, 100), 500, well_conditioned=True, min_diag_value=2.0)

torchsparsegradutils.utils.random_sparse.generate_random_sparse_csr_matrix(size: Tuple[int, int] | Tuple[int, int, int], nnz: int, *, indices_dtype: dtype = torch.int64, values_dtype: dtype = torch.float32, device: device = device(type='cpu'), well_conditioned: bool = False, min_diag_value: float = 1.0) → Tensor[source]

Generate a random sparse CSR matrix.

Creates an unbatched (n_r, n_c) or batched (b, n_r, n_c) CSR tensor with exactly nnz nonzeros per batch item, sampled uniformly at random. Optionally ensures larger diagonal entries for square matrices when well_conditioned=True.

Parameters:

size (tuple of int) – Shape of the output matrix. Must be (n_r, n_c) or (b, n_r, n_c).
nnz (int) – Number of nonzero entries per batch item. Must satisfy 0 <= nnz <= n_r * n_c.
indices_dtype (torch.dtype, default=torch.int64) – Index dtype for the CSR structure (torch.int64 or torch.int32).
values_dtype (torch.dtype, default=torch.float32) – Data type of the nonzero values.
device (torch.device, default=torch.device("cpu")) – Device to allocate the tensors on.
well_conditioned (bool, default=False) – If True and the matrix is square, diagonal entries (if present among the sampled nonzeros) are boosted to be at least min_diag_value.
min_diag_value (float, default=1.0) – Minimum diagonal value when well_conditioned=True.

Returns:

A sparse CSR tensor of shape size with nnz nonzeros per batch (or exactly nnz for the unbatched case).

Return type:

torch.Tensor

Raises:

ValueError – If size has fewer than 2 dims, more than 3 dims, or nnz is out of range.
ValueError – If indices_dtype is not torch.int64 or torch.int32.

Notes

Nonzeros are placed uniformly at random across all entries.

See also

rand_sparse: Convenience dispatcher that selects COO/CSR generator based on layout.
generate_random_sparse_triangular_csr_matrix: Generate triangular CSR matrices.
generate_random_sparse_strictly_triangular_csr_matrix: Generate strictly triangular CSR matrices.

Examples

Unbatched random CSR (80 nonzeros):

>>> A = generate_random_sparse_csr_matrix((20, 30), 80)

Batched random CSR (b=4, each with 120 nonzeros):

>>> A = generate_random_sparse_csr_matrix((4, 50, 40), 120)

Well-conditioned square matrix:

>>> A = generate_random_sparse_csr_matrix((100, 100), 500, well_conditioned=True, min_diag_value=2.0)

torchsparsegradutils.utils.random_sparse.generate_random_sparse_strictly_triangular_coo_matrix(size: Tuple[int, int] | Tuple[int, int, int], nnz: int, *, upper: bool = True, indices_dtype: dtype = torch.int64, values_dtype: dtype = torch.float32, device: device = device(type='cpu'), value_range: Tuple[float, float] = (0.0, 1.0)) → Tensor[source]

Generate a random strictly triangular sparse COO matrix.

Constructs a strictly upper- or lower-triangular sparse matrix in COO format. No diagonal entries are included. Supports unbatched (n, n) or batched (b, n, n) shapes. The number of nonzeros nnz is per batch item (for batched inputs).

Parameters:

size (tuple of int) – Shape of the output matrix. Must be (n, n) for unbatched or (batch_size, n, n) for batched matrices. The matrix must be square.
nnz (int) – Number of nonzero entries per batch. Must satisfy 0 <= nnz <= n*(n-1)/2 for strictly triangular structure.
upper (bool, default=True) – If True, generate strictly upper-triangular coordinates (row < col); otherwise generate strictly lower-triangular coordinates (row > col).
indices_dtype (torch.dtype, default=torch.int64) – Index dtype for the COO indices (torch.int64 or torch.int32).
values_dtype (torch.dtype, default=torch.float32) – Data type of the nonzero values.
device (torch.device, default=torch.device("cpu")) – Device to allocate the tensors on.
value_range (tuple of float, default=(0.0, 1.0)) – Range (min, max) for uniformly sampled nonzero values.

Returns:

A strictly triangular sparse COO tensor of shape size with nnz nonzeros (per batch item if batched).

Return type:

torch.Tensor

Raises:

ValueError – If size has fewer than 2 dims, more than 3 dims, or is not square.
ValueError – If nnz is negative or exceeds n*(n-1)/2.
ValueError – If indices_dtype is not torch.int64 or torch.int32.

Notes

Strictly triangular means the diagonal is excluded by construction.

See also

rand_sparse_tri: Convenience dispatcher for triangular matrices (strict or non-strict; COO/CSR).
generate_random_sparse_triangular_coo_matrix: Generate non-strict triangular COO matrices.

Examples

Unbatched strictly upper-triangular COO (n=100, 300 nonzeros):

>>> A = generate_random_sparse_strictly_triangular_coo_matrix((100, 100), 300, upper=True)

Batched strictly lower-triangular COO (batch=2, n=50, 200 nonzeros per batch):

>>> A = generate_random_sparse_strictly_triangular_coo_matrix((2, 50, 50), 200, upper=False)

torchsparsegradutils.utils.random_sparse.generate_random_sparse_strictly_triangular_csr_matrix(size: Tuple[int, int] | Tuple[int, int, int], nnz: int, *, upper: bool = True, indices_dtype: dtype = torch.int64, values_dtype: dtype = torch.float32, device: device = device(type='cpu'), value_range: Tuple[float, float] = (0.0, 1.0)) → Tensor[source]

Generate a random strictly triangular sparse CSR matrix.

Constructs a strictly upper- or lower-triangular sparse matrix in CSR format. No diagonal entries are included. Supports unbatched (n, n) or batched (b, n, n) shapes. The number of nonzeros nnz is per batch item (for batched inputs).

Parameters:

size (tuple of int) – Shape of the output matrix. Must be (n, n) for unbatched or (batch_size, n, n) for batched matrices. The matrix must be square.
nnz (int) – Number of nonzero entries per batch. Must satisfy 0 <= nnz <= n*(n-1)/2 for strictly triangular structure.
upper (bool, default=True) – If True, generate strictly upper-triangular coordinates (row < col); otherwise generate strictly lower-triangular coordinates (row > col).
indices_dtype (torch.dtype, default=torch.int64) – Index dtype for the CSR structure (torch.int64 or torch.int32).
values_dtype (torch.dtype, default=torch.float32) – Data type of the nonzero values.
device (torch.device, default=torch.device("cpu")) – Device to allocate the tensors on.
value_range (tuple of float, default=(0.0, 1.0)) – Range (min, max) for uniformly sampled nonzero values.

Returns:

A strictly triangular sparse CSR tensor of shape size with nnz nonzeros (per batch item if batched).

Return type:

torch.Tensor

Raises:

ValueError – If size has fewer than 2 dims, more than 3 dims, or is not square.
ValueError – If nnz exceeds n*(n-1)/2.
ValueError – If indices_dtype is not torch.int64 or torch.int32.

Notes

Strictly triangular means the diagonal is excluded by construction.

See also

rand_sparse_tri: Convenience dispatcher for triangular matrices (strict or non-strict; COO/CSR).
generate_random_sparse_triangular_csr_matrix: Generate non-strict triangular CSR matrices.

Examples

Unbatched strictly upper-triangular CSR (n=100, 300 nonzeros):

>>> A = generate_random_sparse_strictly_triangular_csr_matrix((100, 100), 300, upper=True)

Batched strictly lower-triangular CSR (batch=2, n=50, 200 nonzeros per batch):

>>> A = generate_random_sparse_strictly_triangular_csr_matrix((2, 50, 50), 200, upper=False)

torchsparsegradutils.utils.random_sparse.generate_random_sparse_triangular_coo_matrix(size: Tuple[int, int] | Tuple[int, int, int], nnz: int, *, upper: bool = True, indices_dtype: dtype = torch.int64, values_dtype: dtype = torch.float32, device: device = device(type='cpu'), value_range: Tuple[float, float] = (0.0, 1.0), well_conditioned: bool = False, min_diag_value: float = 1.0) → Tensor[source]

Generate a random sparse COO matrix with non-strict triangular structure.

Constructs a triangular sparse matrix in COO format that always includes all diagonal elements. Remaining entries are placed randomly in the upper or lower triangular region depending on upper.

Parameters:

size (tuple of int) – Shape of the matrix. Must be (n, n) for unbatched or (batch_size, n, n) for batched matrices. The matrix must be square.
nnz (int) – Number of nonzero elements per batch. Must satisfy n <= nnz <= n*(n+1)/2. All diagonal elements are included, so nnz must be at least n.
upper (bool, default=True) – If True, generates an upper triangular matrix (row <= col). If False, generates a lower triangular matrix (row >= col).
indices_dtype (torch.dtype, default=torch.int64) – Data type for sparse indices (torch.int64 or torch.int32).
values_dtype (torch.dtype, default=torch.float32) – Data type for nonzero values.
device (torch.device, default=torch.device("cpu")) – Device on which to allocate tensors.
value_range (tuple of float, default=(0.0, 1.0)) – Range (min, max) for random nonzero values.
well_conditioned (bool, default=False) – If True, ensures diagonal elements are large enough for numerical stability.
min_diag_value (float, default=1.0) – Minimum diagonal value if well_conditioned=True.

Returns:

Sparse COO tensor of shape size with triangular structure and diagonal included.

Return type:

torch.Tensor

Raises:

ValueError – If size is not 2D or 3D, not square, or if nnz violates constraints.
ValueError – If indices_dtype is not torch.int64 or torch.int32.

Notes

Always includes all diagonal entries (cannot be excluded).
Off-diagonal nonzeros are chosen randomly within the triangular region.
Use generate_random_sparse_strictly_triangular_coo_matrix for strictly triangular matrices (without diagonal).
Use generate_random_sparse_triangular_csr_matrix for CSR layout instead.

See also

rand_sparse_tri: Convenience dispatcher for triangular matrices (non-strict; COO/CSR).

Examples

Generate a 100×100 lower triangular COO matrix with 300 nonzeros:

>>> A = generate_random_sparse_triangular_coo_matrix((100, 100), 300, upper=False)

Generate a well-conditioned upper triangular COO matrix with minimum diagonal value 2.0:

>>> A = generate_random_sparse_triangular_coo_matrix(
...     (50, 50), 200, upper=True, well_conditioned=True, min_diag_value=2.0
... )

torchsparsegradutils.utils.random_sparse.generate_random_sparse_triangular_csr_matrix(size: Tuple[int, int] | Tuple[int, int, int], nnz: int, *, upper: bool = True, indices_dtype: dtype = torch.int64, values_dtype: dtype = torch.float32, device: device = device(type='cpu'), value_range: Tuple[float, float] = (0.0, 1.0), well_conditioned: bool = False, min_diag_value: float = 1.0) → Tensor[source]

Generate a random sparse CSR matrix with non-strict triangular structure.

Constructs a triangular sparse matrix in CSR format that always includes all diagonal elements. Remaining entries are placed randomly in the upper or lower triangular region depending on upper.

Parameters:

size (tuple of int) – Shape of the matrix. Must be (n, n) for unbatched or (batch_size, n, n) for batched matrices. The matrix must be square.
nnz (int) – Number of nonzero elements per batch. Must satisfy n <= nnz <= n*(n+1)/2. All diagonal elements are included, so nnz must be at least n.
upper (bool, default=True) – If True, generates an upper triangular matrix (row <= col). If False, generates a lower triangular matrix (row >= col).
indices_dtype (torch.dtype, default=torch.int64) – Data type for sparse indices (torch.int64 or torch.int32).
values_dtype (torch.dtype, default=torch.float32) – Data type for nonzero values.
device (torch.device, default=torch.device("cpu")) – Device on which to allocate tensors.
value_range (tuple of float, default=(0.0, 1.0)) – Range (min, max) for random nonzero values.
well_conditioned (bool, default=False) – If True, ensures diagonal elements are large enough for numerical stability.
min_diag_value (float, default=1.0) – Minimum diagonal value if well_conditioned=True.

Returns:

Sparse CSR tensor of shape size with triangular structure and diagonal included.

Return type:

torch.Tensor

Raises:

ValueError – If size is not 2D or 3D, not square, or if nnz violates constraints.
ValueError – If indices_dtype is not torch.int64 or torch.int32.

Notes

Always includes all diagonal entries (cannot be excluded).
Off-diagonal nonzeros are chosen randomly within the triangular region.
Use generate_random_sparse_strictly_triangular_csr_matrix for strictly triangular matrices (without diagonal).
Use generate_random_sparse_triangular_coo_matrix for COO layout instead.

See also

rand_sparse_tri: Convenience dispatcher for triangular matrices (non-strict; COO/CSR).

Examples

Generate a 100×100 upper triangular CSR matrix with 300 nonzeros:

>>> A = generate_random_sparse_triangular_csr_matrix((100, 100), 300, upper=True)

Generate a well-conditioned lower triangular CSR matrix with min diagonal value 2.0:

>>> A = generate_random_sparse_triangular_csr_matrix(
...     (50, 50), 200, upper=False, well_conditioned=True, min_diag_value=2.0
... )

Statistical Distribution Helpers

Statistical tests for validating multivariate distributions using confidence regions.

This module provides statistical tests for validating multivariate normal distributions by testing whether observed sample statistics are consistent with hypothesized population parameters using confidence regions rather than traditional hypothesis testing.

Notes

Understanding Confidence Regions

These tests use confidence regions to determine whether observed statistics are consistent with hypothesized parameters:

A confidence region with level C (e.g., 0.95) means that if we repeated the experiment many times, C% of the confidence regions would contain the true parameter values
HIGHER confidence levels (closer to 1.0) create LARGER regions, making tests EASIER to pass
LOWER confidence levels create SMALLER regions, making tests HARDER to pass

Examples of confidence levels:

confidence_level = 0.50 → 50% confidence region (very strict)
confidence_level = 0.95 → 95% confidence region (moderately strict)
confidence_level = 0.99 → 99% confidence region (more lenient)

When to use each test:

Use higher confidence levels (0.95-0.99) for practical distribution validation
Use moderate levels (0.80-0.95) for unit testing with some tolerance
Use lower levels (0.50-0.80) for strict validation or when you expect the parameters to be very close

Practical interpretation:

Think of confidence levels as “how generous should I be?”:

0.99: “I’ll accept the hypothesis unless there’s very strong evidence against it”
0.95: “I’ll accept the hypothesis unless there’s moderate evidence against it”
0.50: “I’ll only accept the hypothesis if it’s very close to the data”

This is the OPPOSITE of traditional hypothesis testing where smaller p-values (like 0.05) indicate stronger evidence against the null hypothesis. Here, larger confidence levels create more permissive acceptance regions.

For implementation details, see the references for Hotelling’s test [1b] and Nagao’s test [2b].

References

[1b]

Hotelling, H. (1947). Multivariate Quality Control. In C. Eisenhart, M. W. Hastay, and W. A. Wallis, eds. Techniques of Statistical Analysis. New York: McGraw-Hill.

[2b]

Nagao, H. (1973). On Some Test Criteria for Covariance Matrix. The Annals of Statistics, Vol. 1, No. 4, pp. 700-709.

torchsparsegradutils.utils.dist_stats_helpers.mean_hotelling_t2_test(sample_mean: Tensor, true_mean: Tensor, sample_cov: Tensor, n: int, confidence_level: float = 0.95) → Tuple[Tensor, Tensor, float][source]

One-sample Hotelling T² test for multivariate mean equality using confidence regions.

This test checks whether the hypothesized mean \(\boldsymbol{\mu}_0\) lies inside the joint confidence region around the sample mean \(\bar{\boldsymbol{x}}\), accounting for covariance among variables. The confidence region is an ellipsoid in \(p\)-dimensional space.

Test statistic (Hotelling’s T²):

\[T^2 \;=\; n\, (\bar{\boldsymbol{x}} - \boldsymbol{\mu}_0)^{\top} \hat{\boldsymbol{\Sigma}}^{-1} (\bar{\boldsymbol{x}} - \boldsymbol{\mu}_0)\]

with the distributional relationship (for unknown covariance):

\[T^2 \;\sim\; \frac{p(n-1)}{n-p}\,F_{p,\,n-p}.\]

Acceptance criterion (confidence region): \(\boldsymbol{\mu}_0\) is accepted at the specified confidence level if

\[T^2 \;\le\; \frac{p(n-1)}{n-p} \; F_{p,\,n-p;\,\text{confidence level}},\]

equivalently

\[(\bar{\boldsymbol{x}} - \boldsymbol{\mu})^{\top} \hat{\boldsymbol{\Sigma}}^{-1} (\bar{\boldsymbol{x}} - \boldsymbol{\mu}) \;\le\; \frac{p(n-1)}{n(n-p)} \; F_{p,\,n-p;\,\text{confidence level}},\]

for all \(\boldsymbol{\mu}\) in the confidence region.

Parameters:

sample_mean (torch.Tensor, shape (B, p)) – Sample mean vector where B is the batch size and p is the dimension.
true_mean (torch.Tensor, shape (B, p)) – Hypothesized true mean vector.
sample_cov (torch.Tensor, shape (B, p, p)) – Sample covariance matrix.
n (int) – Sample size used to compute the sample mean and covariance.
confidence_level (float, default=0.95) – Confidence level for the region. Higher values create larger regions (easier to pass). Must be in (0, 1).

Returns:

test_result (torch.Tensor, shape (B,), dtype=bool) – Boolean tensor indicating whether the hypothesized mean lies within the confidence region (True) or outside it (False).
t2_statistic (torch.Tensor, shape (B,), dtype=float) – T² test statistic values.
t2_threshold (float) – T² threshold value for the confidence region.

Notes

Interpretation of confidence levels:

Higher confidence levels (e.g., 0.99) create larger regions, making the test more lenient (easier to accept the hypothesis)
Lower confidence levels (e.g., 0.50) create smaller regions, making the test more strict (harder to accept the hypothesis)

For practical distribution validation, use confidence levels of 0.95-0.99. For strict unit testing, consider 0.80-0.95.

This implementation follows standard Hotelling’s T² methodology (see e.g., Hotelling’s T-squared distribution [1c] [2c]). For confidence regions, see Confidence region [3c].

References

[1c]

Wikipedia: Hotelling’s T-squared distribution. https://en.wikipedia.org/wiki/Hotelling%27s_T-squared_distribution

[2c]

Hotelling, H. (1947). Multivariate Quality Control. In C. Eisenhart, M. W. Hastay, and W. A. Wallis, eds. Techniques of Statistical Analysis. New York: McGraw-Hill.

[3c]

Wikipedia: Confidence region. https://en.wikipedia.org/wiki/Confidence_region

Examples

Test if sample statistics are consistent with known parameters:

>>> import torch
>>> from torch.distributions import MultivariateNormal
>>> from torchsparsegradutils.utils import mean_hotelling_t2_test
>>>
>>> # Set random seed for reproducible results
>>> _ = torch.manual_seed(42)
>>>
>>> # True parameters
>>> true_mean = torch.tensor([[0.0, 0.0]])    # shape (1, 2)
>>> true_cov = torch.eye(2).unsqueeze(0)      # shape (1, 2, 2)
>>> n = 1000
>>>
>>> # Generate sample data from the true distribution
>>> dist = MultivariateNormal(true_mean.squeeze(0), true_cov.squeeze(0))
>>> samples = dist.sample((n,)).unsqueeze(1)  # Add batch dimension
>>> sample_mean = samples.mean(0)
>>> sample_cov = torch.cov(samples.squeeze(1).T).unsqueeze(0)
>>>
>>> # Test with 95% confidence
>>> result, t2_stat, threshold = mean_hotelling_t2_test(
...     sample_mean, true_mean, sample_cov, n, confidence_level=0.95
... )
>>> print(f"Test passed: {result.item()}")  # Should be True
Test passed: True
>>> print(f"T² statistic: {t2_stat.item():.3f}, threshold: {threshold:.3f}")
T² statistic: ..., threshold: ...

torchsparsegradutils.utils.dist_stats_helpers.cov_nagao_test(emp_cov: Tensor, ref_cov: Tensor, n: int, confidence_level: float = 0.95) → Tuple[Tensor, Tensor, float][source]

Nagao’s (1973) one-sample test for covariance matrix equality using confidence regions.

This test determines whether the hypothesized covariance matrix \(\boldsymbol{\Sigma}_0\) is consistent with the empirical covariance matrix, by checking if the test statistic lies within the appropriate confidence region under the \(\chi^2\) distribution.

The test standardizes the empirical covariance matrix:

\[\mathbf{W} = \boldsymbol{\Sigma}_0^{-1/2} \, \hat{\boldsymbol{\Sigma}} \, \boldsymbol{\Sigma}_0^{-1/2}\]

Here, \(\hat{\boldsymbol{\Sigma}}\) denotes the empirical (sample) covariance, and \(\mathbf{W}\) is the covariance whitened with respect to \(\boldsymbol{\Sigma}_0\) (i.e., a similarity transform that equals the identity when \(\hat{\boldsymbol{\Sigma}} = \boldsymbol{\Sigma}_0\)). The factor \(\boldsymbol{\Sigma}_0^{-1/2}\) is computed from the Cholesky decomposition of \(\boldsymbol{\Sigma}_0\). Under \(H_0: \boldsymbol{\Sigma} = \boldsymbol{\Sigma}_0\), the matrix \(\mathbf{W}\) should be close to the identity matrix \(\mathbf{I}\).

The test statistic is:

\[T_N = \frac{n}{2} \cdot \left\|\mathbf{W} - \mathbf{I}\right\|_F^2\]

Where \(\|\cdot\|_F\) is the Frobenius norm.

Distribution under \(H_0\) (asymptotically):

\[T_N \;\sim\; \chi^2_{\,\nu}, \qquad \nu = \tfrac{p(p+1)}{2},\]

where \(p\) is the dimension.

Acceptance criterion (confidence region): at the specified confidence level, accept \(\boldsymbol{\Sigma}_0\) if

\[T_N \;\leq\; \chi^2_{\nu;\,\text{confidence level}} \quad \text{with} \quad \nu = \tfrac{p(p+1)}{2}.\]

Parameters:

emp_cov (torch.Tensor, shape (B, p, p)) – Empirical (sample) covariance matrix where B is the batch size and p is the dimension.
ref_cov (torch.Tensor, shape (B, p, p)) – Reference (hypothesized) covariance matrix.
n (int) – Sample size used to compute the empirical covariance matrix.
confidence_level (float, default=0.95) – Confidence level for the region. Higher values create larger regions (easier to pass). Must be in (0, 1).

Returns:

test_result (torch.Tensor, shape (B,), dtype=bool) – Boolean tensor indicating whether the hypothesized covariance lies within the confidence region (True) or outside it (False).
t_n_statistic (torch.Tensor, shape (B,), dtype=float) – T_N test statistic values.
chi2_threshold (float) – \(\chi^2\) threshold value for the confidence region.

Notes

This test is often more stable than raw Frobenius norm tests when the reference covariance \(\boldsymbol{\Sigma}_0\) is ill-conditioned.

Interpretation of confidence levels:

Higher confidence levels (e.g., 0.99) create larger regions, making the test more lenient (easier to accept the hypothesis)
Lower confidence levels (e.g., 0.50) create smaller regions, making the test more strict (harder to accept the hypothesis)

For practical distribution validation, use confidence levels of 0.95-0.99. For strict unit testing, consider 0.80-0.95.

This implementation follows Nagao’s test methodology [1d].

References

[1d]

Nagao, H. (1973). On Some Test Criteria for Covariance Matrix. The Annals of Statistics, Vol. 1, No. 4, pp. 700-709. Institute of Mathematical Statistics. Stable URL: https://www.jstor.org/stable/2958313 Equation 3.3

Examples

Test if sample covariance is consistent with known covariance:

>>> import torch
>>> from torch.distributions import MultivariateNormal
>>> from torchsparsegradutils.utils import cov_nagao_test
>>>
>>> # Set random seed for reproducible results
>>> _ = torch.manual_seed(42)
>>>
>>> # True parameters
>>> true_mean = torch.tensor([[0.0, 0.0]])    # shape (1, 2)
>>> ref_cov = torch.tensor([[[1.0, 0.0], [0.0, 1.0]]])  # shape (1, 2, 2)
>>> n = 1000
>>>
>>> # Generate sample data from the true distribution
>>> dist = MultivariateNormal(true_mean.squeeze(0), ref_cov.squeeze(0))
>>> samples = dist.sample((n,)).unsqueeze(1)  # Add batch dimension
>>> emp_cov = torch.cov(samples.squeeze(1).T).unsqueeze(0)
>>>
>>> # Test with 95% confidence
>>> result, t_n_stat, threshold = cov_nagao_test(
...     emp_cov, ref_cov, n, confidence_level=0.95
... )
>>> print(f"Test passed: {result.item()}")  # Should be True
Test passed: True
>>> print(f"T_N statistic: {t_n_stat.item():.3f}, threshold: {threshold:.3f}")
T_N statistic: ..., threshold: ...