trnsparse 0.3.1__tar.gz → 0.3.2__tar.gz

{trnsparse-0.3.1 → trnsparse-0.3.2}/CHANGELOG.md

@@ -5,6 +5,35 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.3.2] — 2026-04-14
+
+### Added
+
+- **`cg_bsr`** and **`power_iteration_bsr`** — Conjugate Gradient and
+  power iteration on block-sparse row matrices. Plumbing on top of
+  `bsr_spmm` (one kernel dispatch per iteration). Closes Phase 1 of
+  #22 on-chip iterative solvers.
+- **`jacobi_preconditioner_bsr(A)`** — builds a diagonal preconditioner
+  for `cg_bsr`'s `M=` argument.
+- **`bsr_diagonal(A)`** — extracts the main diagonal from a BSR matrix.
+- **`docs/iterative_solvers.md`** — design note covering the v0.3.2
+  plumbing and the v0.4.0 fused-kernel goal (#24). Explains the
+  architectural win Trainium offers (A SBUF-resident across iterations)
+  vs the current per-iteration HBM round-trip.
+- **`tests/test_iterative.py`** — 8 CPU tests including scipy parity at
+  `atol=1e-4` on a 128×128 SPD system.
+- **`benchmarks/bench_iterative.py`** — cg_bsr vs scipy.sparse.linalg.cg.
+  At 128×128 SPD: scipy 310 μs, trnsparse 369 μs (1.19×).
+
+### Notes
+
+- Algorithm body for CG is a local copy of `trnsolver.iterative.cg`;
+  kept local to avoid a cross-repo runtime dependency for one function.
+- v0.4.0 will layer the fused CG/power-iteration NKI kernel on top —
+  tracked in #24. The API stays stable across the transition; users
+  upgrading from v0.3.2 get the fused-kernel speedup automatically
+  when the fused path is available.
+
 ## [0.3.1] — 2026-04-14
 
 ### Changed

{trnsparse-0.3.1/trnsparse.egg-info → trnsparse-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trnsparse
-Version: 0.3.1
+Version: 0.3.2
 Summary: Sparse matrix operations for AWS Trainium via NKI
 Author-email: Scott Friedman <scttfrdmn@gmail.com>
 License-Expression: Apache-2.0

trnsparse-0.3.2/benchmarks/bench_iterative.py

@@ -0,0 +1,50 @@
+"""Iterative-solver benchmarks: trnsparse.cg_bsr vs scipy baseline.
+
+The v0.3.2 plumbing dispatches one `bsr_spmm` call per CG iteration.
+On NKI that means one kernel launch + HBM round-trip per iteration.
+Expect the current path to be dominated by dispatch overhead compared
+to scipy's compiled C loop, which motivates the v0.4.0 fused-kernel
+follow-up.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+import trnsparse
+
+
+@pytest.fixture(params=[128, 256])
+def iter_size(request):
+    return request.param
+
+
+@pytest.fixture
+def spd_bsr_and_dense(iter_size):
+    torch.manual_seed(0)
+    n = iter_size
+    M = torch.randn(n, n)
+    A_dense = M @ M.T + n * torch.eye(n)
+    A_bsr = trnsparse.BSRMatrix.from_dense(A_dense, block_size=128)
+    b = torch.randn(n)
+    return A_dense, A_bsr, b
+
+
+def test_cg_bsr_trnsparse(benchmark, spd_bsr_and_dense):
+    _, A_bsr, b = spd_bsr_and_dense
+    benchmark(lambda: trnsparse.cg_bsr(A_bsr, b, tol=1e-8, max_iter=2 * b.shape[0]))
+
+
+def test_cg_scipy(benchmark, spd_bsr_and_dense):
+    sp = pytest.importorskip("scipy.sparse")
+    spla = pytest.importorskip("scipy.sparse.linalg")
+    A_dense, _, b = spd_bsr_and_dense
+    A_scipy = sp.csr_matrix(A_dense.numpy())
+    b_np = b.numpy()
+    benchmark(lambda: spla.cg(A_scipy, b_np, rtol=1e-8, maxiter=2 * len(b_np)))
+
+
+def test_power_iteration_trnsparse(benchmark, spd_bsr_and_dense):
+    _, A_bsr, _ = spd_bsr_and_dense
+    benchmark(lambda: trnsparse.power_iteration_bsr(A_bsr, max_iter=500, tol=1e-9))

trnsparse-0.3.2/docs/iterative_solvers.md

@@ -0,0 +1,115 @@
+# Iterative solvers over BSR
+
+trnsparse v0.3.2 adds `cg_bsr` and `power_iteration_bsr` — Conjugate
+Gradient and power iteration on block-sparse row matrices. The API is
+stable; the architectural story below explains why there's a v0.4.0
+follow-up.
+
+## Why this matters
+
+Large SPD linear systems and dominant-eigenpair problems show up
+across scientific computing:
+
+- **Quantum chemistry**: Hamiltonian eigenvalue problems (HF, DFT,
+  CI), response equations (CPSCF).
+- **PDE discretizations**: stiffness-matrix solves for finite element
+  methods, graph Laplacian systems.
+- **Graph learning**: spectral embeddings, PageRank-like iterations.
+
+The matrix `A` in each case is typically block-sparse (Fock matrices
+after Schwarz screening; FEM stiffness tied to mesh connectivity;
+graph adjacency). BSR is the Trainium-native representation for those
+matrices (see `architecture.md`).
+
+## v0.3.2 — plumbing
+
+```python
+import trnsparse
+
+A = trnsparse.BSRMatrix.from_dense(fock_matrix, block_size=128)
+b = compute_rhs()
+
+x, iters, rel = trnsparse.cg_bsr(A, b, tol=1e-6, max_iter=1000)
+# Jacobi-preconditioned variant:
+M = trnsparse.jacobi_preconditioner_bsr(A)
+x, iters, rel = trnsparse.cg_bsr(A, b, tol=1e-6, M=M)
+
+lam, v, iters = trnsparse.power_iteration_bsr(A, max_iter=500)
+```
+
+Under the hood, each CG iteration calls `bsr_spmm(A, x.unsqueeze(1))`
+once. On the NKI backend that's one kernel dispatch + one HBM
+round-trip per iteration. On CPU it's `torch.sparse`-backed and
+roughly on par with `scipy.sparse.linalg.cg` (benchmarked: 369 μs vs
+310 μs at 128×128 SPD, 1.19× slower).
+
+## v0.4.0 — fused kernel with SBUF-resident A
+
+The architectural claim from
+[#22](https://github.com/trnsci/trnsparse/issues/22): Trainium's 32 GB
+SBUF per NeuronCore fits a 5000×5000 BSR Hamiltonian on-chip. CG
+doesn't need to round-trip `A` to HBM at all — only `x`, `r`, and `p`.
+
+The shape of the v0.4.0 kernel:
+
+```python
+@nki.jit
+def _cg_spd_kernel(A_blocks, A_cols, A_row_ptrs, b, max_iter):
+    # Load A blocks once into SBUF at the top.
+    A_sbuf = nl.load(A_blocks)
+
+    # State: x, r, p in SBUF registers.
+    x = nl.zeros(...)
+    r = nl.copy(b)
+    p = nl.copy(r)
+    rr = nl.reduce(r * r)
+
+    for k in nl.affine_range(max_iter):
+        Ap = _bsr_matvec_sbuf(A_sbuf, A_cols, A_row_ptrs, p)   # all SBUF
+        pAp = nl.reduce(p * Ap)
+        alpha = rr / pAp
+        x = x + alpha * p
+        r = r - alpha * Ap
+        rr_new = nl.reduce(r * r)
+        beta = rr_new / rr
+        p = r + beta * p
+        rr = rr_new
+
+    return x, residual_norm_history
+```
+
+Fixed `max_iter`, no early exit — returning the full residual history
+lets the host pick the convergence point post-hoc (standard
+NKI-inside-loop constraint: no dynamic control flow).
+
+**Expected regime where this wins**: when
+`max_iter × dispatch_overhead > fused_kernel_cost + hbm_load_A_once`.
+For a 4000×4000 BSR Hamiltonian with ~100 iterations to convergence,
+that's roughly an order of magnitude of wall-time reduction.
+
+**What needs to land first**:
+
+1. Simulator-iterated kernel skeleton (now tractable thanks to the
+   `nki-simulator` CI gate — see the NKI 0.3.0 migration in v0.3.1).
+2. SBUF sizing model for `A` — some workloads overflow 32 GB; then
+   fall back to the v0.3.2 plumbing.
+3. Dispatcher logic in `cg_bsr` that picks the fused kernel when
+   `max_iter * n` exceeds a threshold.
+
+Tracked in a dedicated sub-issue on `trnsci/trnsparse`.
+
+## Why CG and power iteration, not GMRES / Lanczos / Davidson
+
+v0.3.2 covers the two most common algorithm families:
+
+- **CG** for SPD systems — the workhorse for Hamiltonian solves,
+  stiffness systems, and many PDE discretizations.
+- **Power iteration** for dominant eigenpairs — the starting point
+  for spectral methods, PageRank-like fixed points, and the iteration
+  core inside Lanczos / Arnoldi / Davidson.
+
+GMRES (general non-symmetric systems), Lanczos / Arnoldi (full
+spectrum), and Davidson (interior eigenvalues) are follow-ups when
+users ask for them. The v0.4.0 fused kernel's structure (load A once,
+iterate on-chip) generalizes trivially to those variants — it's the
+same architectural pattern.

{trnsparse-0.3.1 → trnsparse-0.3.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "trnsparse"
-version = "0.3.1"
+version = "0.3.2"
 description = "Sparse matrix operations for AWS Trainium via NKI"
 readme = "README.md"
 license = "Apache-2.0"

trnsparse-0.3.2/tests/test_iterative.py

@@ -0,0 +1,135 @@
+"""Iterative solver tests on BSR matrices (CPU plumbing — #22 Phase 1)."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+import trnsparse
+
+
+def _random_spd_bsr(
+    n: int, block_size: int = 128, seed: int = 0
+) -> tuple[torch.Tensor, trnsparse.BSRMatrix]:
+    """Build a dense SPD matrix of size (n, n), return (dense, BSR view).
+
+    `A = M @ M.T + n*I` guarantees SPD. Block-storing a dense SPD matrix
+    is valid — BSR just chooses which blocks to keep; SPD-ness is a
+    per-element property that survives the block partition.
+    """
+    torch.manual_seed(seed)
+    M = torch.randn(n, n)
+    A_dense = M @ M.T + n * torch.eye(n)
+    return A_dense, trnsparse.BSRMatrix.from_dense(A_dense, block_size=block_size)
+
+
+class TestCgBsr:
+    def test_identity_converges_immediately(self):
+        """A = I → x = b in one iteration (or zero for zero RHS)."""
+        n = 128
+        I_dense = torch.eye(n)
+        I_bsr = trnsparse.BSRMatrix.from_dense(I_dense, block_size=128)
+        b = torch.randn(n)
+
+        x, iters, rel = trnsparse.cg_bsr(I_bsr, b, tol=1e-8, max_iter=50)
+        torch.testing.assert_close(x, b, atol=1e-5, rtol=1e-5)
+        assert iters <= 2, f"identity should converge in 1 iter, took {iters}"
+        assert rel < 1e-8
+
+    def test_parity_vs_scipy(self):
+        """Solution from cg_bsr matches scipy.sparse.linalg.cg within tol."""
+        sp = pytest.importorskip("scipy.sparse")
+        spla = pytest.importorskip("scipy.sparse.linalg")
+
+        n = 128
+        A_dense, A_bsr = _random_spd_bsr(n, block_size=128, seed=1)
+        b = torch.randn(n)
+
+        x_ours, iters_ours, rel_ours = trnsparse.cg_bsr(A_bsr, b, tol=1e-8, max_iter=2 * n)
+        # scipy's cg takes rtol (relative to ||b||) by default in modern scipy
+        A_scipy = sp.csr_matrix(A_dense.numpy())
+        x_scipy, info = spla.cg(A_scipy, b.numpy(), rtol=1e-8, maxiter=2 * n)
+
+        assert info == 0, "scipy CG should converge on this well-conditioned SPD system"
+        np.testing.assert_allclose(x_ours.numpy(), x_scipy, atol=1e-4, rtol=1e-4)
+        assert rel_ours < 1e-7
+
+    def test_jacobi_preconditioner_fewer_iters(self):
+        """A strongly diagonally-dominant system should converge faster
+        with Jacobi preconditioning than without.
+        """
+        torch.manual_seed(2)
+        n = 256
+        # Diagonally dominant: large diagonal, small off-diag.
+        A_dense = 0.01 * torch.randn(n, n)
+        A_dense = 0.5 * (A_dense + A_dense.T)  # symmetric
+        A_dense = A_dense + torch.diag(1.0 + 10.0 * torch.rand(n))  # SPD
+        A_bsr = trnsparse.BSRMatrix.from_dense(A_dense, block_size=128)
+        b = torch.randn(n)
+
+        _, iters_plain, _ = trnsparse.cg_bsr(A_bsr, b, tol=1e-8, max_iter=500)
+        M = trnsparse.jacobi_preconditioner_bsr(A_bsr)
+        _, iters_precond, _ = trnsparse.cg_bsr(A_bsr, b, tol=1e-8, max_iter=500, M=M)
+
+        assert iters_precond <= iters_plain, (
+            f"Jacobi preconditioner should not make it worse "
+            f"(plain={iters_plain}, precond={iters_precond})"
+        )
+
+    def test_zero_rhs_returns_zero(self):
+        """b = 0 → x = 0 in 0 iterations."""
+        _, A_bsr = _random_spd_bsr(128, block_size=128, seed=3)
+        b = torch.zeros(128)
+        x, iters, rel = trnsparse.cg_bsr(A_bsr, b)
+        torch.testing.assert_close(x, b)
+        assert iters == 0
+        assert rel == 0.0
+
+
+class TestPowerIterationBsr:
+    def test_dominant_eigenvalue_diagonal(self):
+        """Diagonal matrix: dominant eigenvalue is max(diag); eigenvector
+        is the corresponding basis vector.
+        """
+        n = 128
+        diag = torch.linspace(0.5, 10.0, n)  # strictly increasing; max = 10
+        A_dense = torch.diag(diag)
+        A_bsr = trnsparse.BSRMatrix.from_dense(A_dense, block_size=128)
+
+        lam, v, iters = trnsparse.power_iteration_bsr(A_bsr, max_iter=500, tol=1e-10)
+        assert abs(lam - 10.0) < 1e-3, f"expected ~10.0, got {lam}"
+        # Eigenvector aligns with e_{n-1} (corresponding to max diag entry).
+        assert abs(abs(v[-1].item()) - 1.0) < 1e-3
+
+    def test_dominant_eigenvalue_vs_torch(self):
+        """Match torch.linalg.eigvalsh's largest eigenvalue within 1e-3."""
+        torch.manual_seed(4)
+        n = 128
+        A_dense, A_bsr = _random_spd_bsr(n, block_size=128, seed=4)
+
+        lam_ours, _, _ = trnsparse.power_iteration_bsr(A_bsr, max_iter=2000, tol=1e-10)
+        lam_ref = torch.linalg.eigvalsh(A_dense).max().item()
+
+        # Power iteration on a random SPD matrix has a spectral gap;
+        # 1e-3 relative tolerance is generous.
+        assert abs(lam_ours - lam_ref) / lam_ref < 1e-3, f"ours={lam_ours}, ref={lam_ref}"
+
+
+class TestBsrDiagonal:
+    def test_diagonal_of_identity(self):
+        I_bsr = trnsparse.BSRMatrix.from_dense(torch.eye(128), block_size=128)
+        d = trnsparse.bsr_diagonal(I_bsr)
+        torch.testing.assert_close(d, torch.ones(128))
+
+    def test_diagonal_of_off_diagonal_blocks_only(self):
+        """A matrix whose only stored blocks are off-diagonal returns
+        zeros on the diagonal.
+        """
+        torch.manual_seed(5)
+        A_dense = torch.zeros(256, 256)
+        A_dense[:128, 128:] = torch.randn(128, 128)
+        A_dense[128:, :128] = torch.randn(128, 128)
+        A_bsr = trnsparse.BSRMatrix.from_dense(A_dense, block_size=128)
+        d = trnsparse.bsr_diagonal(A_bsr)
+        torch.testing.assert_close(d, torch.zeros(256))

{trnsparse-0.3.1 → trnsparse-0.3.2}/trnsparse/__init__.py

@@ -5,9 +5,10 @@ CSR/COO formats, SpMV, SpMM, and integral screening for
 sparse scientific computing. Part of the trnsci scientific computing suite.
 """
 
-__version__ = "0.3.1"
+__version__ = "0.3.2"
 
 from .formats import BSRMatrix, COOMatrix, CSRMatrix, eye_sparse, from_dense, from_scipy
+from .iterative import bsr_diagonal, cg_bsr, jacobi_preconditioner_bsr, power_iteration_bsr
 from .nki import HAS_NKI, get_backend, set_backend
 from .ops import (
     bsr_spmm,
@@ -40,6 +41,10 @@ __all__ = [
     "screen_quartets",
     "density_screen",
     "sparsity_stats",
+    "cg_bsr",
+    "power_iteration_bsr",
+    "jacobi_preconditioner_bsr",
+    "bsr_diagonal",
     "HAS_NKI",
     "set_backend",
     "get_backend",

trnsparse-0.3.2/trnsparse/iterative.py

@@ -0,0 +1,190 @@
+"""Iterative linear solvers for BSR sparse matrices.
+
+v0.3.2 — Phase 1 plumbing (#22). Python-level CG and power iteration
+on top of `bsr_spmm` as the matvec. A gets reloaded from HBM on every
+iteration; v0.4.0 (follow-up) delivers the SBUF-resident fused NKI
+kernel that keeps A on-chip across iterations — the architectural win.
+
+Algorithm body for CG mirrors `trnsolver.iterative.cg` at commit-time,
+kept local to avoid a trnsolver runtime dep.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+
+from .formats import BSRMatrix
+from .ops import bsr_spmm
+
+
+def bsr_diagonal(A: BSRMatrix) -> torch.Tensor:
+    """Extract the main diagonal of a BSR matrix as a dense vector.
+
+    The diagonal lives inside the diagonal blocks (block row i at block
+    column i). Blocks that aren't diagonal contribute nothing. Returns
+    a zero vector if no diagonal block is stored (a fully zero-diagonal
+    matrix would be degenerate as a preconditioner anyway).
+    """
+    m, _ = A.shape
+    b = A.block_size
+    M_tiles = (m + b - 1) // b
+    out = torch.zeros(M_tiles * b, dtype=A.dtype)
+    for i in range(M_tiles):
+        start = A.block_row_ptrs[i].item()
+        end = A.block_row_ptrs[i + 1].item()
+        for k in range(start, end):
+            if A.block_col_indices[k].item() == i:
+                out[i * b : (i + 1) * b] = torch.diagonal(A.blocks[k])
+                break
+    return out[:m]
+
+
+def _bsr_matvec(A: BSRMatrix) -> Callable[[torch.Tensor], torch.Tensor]:
+    """Wrap `bsr_spmm` so CG / power iteration can pass it a vector."""
+
+    def matvec(x: torch.Tensor) -> torch.Tensor:
+        return bsr_spmm(A, x.unsqueeze(1)).squeeze(1)
+
+    return matvec
+
+
+def cg_bsr(
+    A: BSRMatrix,
+    b: torch.Tensor,
+    x0: torch.Tensor | None = None,
+    tol: float = 1e-6,
+    max_iter: int = 1000,
+    M: Callable[[torch.Tensor], torch.Tensor] | torch.Tensor | None = None,
+) -> tuple[torch.Tensor, int, float]:
+    """Conjugate Gradient solver for `A @ x = b` with A SPD and BSR-stored.
+
+    Thin specialization of the standard CG loop: the matvec goes through
+    `bsr_spmm`, so on the NKI backend each iteration dispatches one
+    kernel call. That amortizes poorly across many iterations — the
+    v0.4.0 fused kernel fuses the loop and keeps A SBUF-resident.
+
+    Args:
+        A: SPD BSR matrix (N, N).
+        b: Right-hand side, shape (N,).
+        x0: Initial guess (default: zeros).
+        tol: Relative-residual tolerance on `||r|| / ||b||`.
+        max_iter: Iteration cap.
+        M: Preconditioner — callable `M(r) → M^{-1} r` or a dense
+            inverse tensor. Build a Jacobi preconditioner from
+            `bsr_diagonal(A)` for the common diagonally-dominant case.
+
+    Returns:
+        `(x, iters, rel_residual)` — solution, iteration count, final
+        relative residual.
+    """
+    n = b.shape[0]
+    matvec = _bsr_matvec(A)
+
+    if M is None:
+        precond: Callable[[torch.Tensor], torch.Tensor] | None = None
+    elif callable(M):
+        precond = M
+    else:
+        M_tensor = M
+
+        def precond(r: torch.Tensor) -> torch.Tensor:
+            return torch.mv(M_tensor, r)
+
+    x = x0.clone() if x0 is not None else torch.zeros(n, dtype=b.dtype, device=b.device)
+    r = b - matvec(x)
+    b_norm = torch.linalg.norm(b).item()
+    if b_norm < 1e-15:
+        return x, 0, 0.0
+
+    z = precond(r) if precond is not None else r.clone()
+    p = z.clone()
+    rz = torch.dot(r, z).item()
+
+    for k in range(max_iter):
+        Ap = matvec(p)
+        pAp = torch.dot(p, Ap).item()
+        if abs(pAp) < 1e-30:
+            break
+
+        alpha = rz / pAp
+        x = x + alpha * p
+        r = r - alpha * Ap
+
+        r_norm = torch.linalg.norm(r).item()
+        rel = r_norm / b_norm
+        if rel < tol:
+            return x, k + 1, rel
+
+        z = precond(r) if precond is not None else r
+        rz_new = torch.dot(r, z).item()
+        beta = rz_new / rz
+        p = z + beta * p
+        rz = rz_new
+
+    return x, max_iter, r_norm / b_norm
+
+
+def jacobi_preconditioner_bsr(
+    A: BSRMatrix,
+) -> Callable[[torch.Tensor], torch.Tensor]:
+    """Build a Jacobi (diagonal) preconditioner for a BSR matrix.
+
+    Returns a callable `M(r) = r / diag(A)`. Cheap + effective when A
+    is diagonally dominant.
+    """
+    d = bsr_diagonal(A)
+    if torch.any(d.abs() < 1e-15):
+        raise ValueError("Jacobi preconditioner: diagonal has near-zero entries")
+    inv_d = 1.0 / d
+
+    def precond(r: torch.Tensor) -> torch.Tensor:
+        return r * inv_d
+
+    return precond
+
+
+def power_iteration_bsr(
+    A: BSRMatrix,
+    v0: torch.Tensor | None = None,
+    max_iter: int = 100,
+    tol: float = 1e-9,
+) -> tuple[float, torch.Tensor, int]:
+    """Find the dominant eigenpair (λ, v) of a BSR matrix via power iteration.
+
+    Converges to the eigenvector with largest `|λ|`. Tolerance is on
+    the Rayleigh-quotient change between iterations.
+
+    Args:
+        A: Square BSR matrix (N, N). Doesn't need to be SPD — power
+            iteration works for any matrix with a unique dominant
+            eigenvalue, but convergence slows as the spectral gap
+            shrinks.
+        v0: Initial vector (default: random unit vector).
+        max_iter: Iteration cap.
+        tol: Convergence on `|λ_k - λ_{k-1}|`.
+
+    Returns:
+        `(eigenvalue, eigenvector, iters)`.
+    """
+    n = A.shape[0]
+    matvec = _bsr_matvec(A)
+
+    v = v0.clone() if v0 is not None else torch.randn(n, dtype=A.dtype)
+    v = v / torch.linalg.norm(v)
+
+    lam_prev = 0.0
+    for k in range(max_iter):
+        Av = matvec(v)
+        lam = torch.dot(v, Av).item()
+        norm = torch.linalg.norm(Av).item()
+        if norm < 1e-30:
+            return 0.0, v, k + 1
+        v = Av / norm
+
+        if abs(lam - lam_prev) < tol:
+            return lam, v, k + 1
+        lam_prev = lam
+
+    return lam, v, max_iter

{trnsparse-0.3.1 → trnsparse-0.3.2/trnsparse.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trnsparse
-Version: 0.3.1
+Version: 0.3.2
 Summary: Sparse matrix operations for AWS Trainium via NKI
 Author-email: Scott Friedman <scttfrdmn@gmail.com>
 License-Expression: Apache-2.0

{trnsparse-0.3.1 → trnsparse-0.3.2}/trnsparse.egg-info/SOURCES.txt

@@ -12,6 +12,7 @@ pyproject.toml
 .github/workflows/notify-umbrella.yml
 .github/workflows/publish.yml
 benchmarks/bench_bsr_spmm.py
+benchmarks/bench_iterative.py
 benchmarks/bench_screening.py
 benchmarks/bench_spmm.py
 benchmarks/bench_spmv.py
@@ -22,6 +23,7 @@ docs/aws_setup.md
 docs/benchmarks.md
 docs/index.md
 docs/installation.md
+docs/iterative_solvers.md
 docs/migration_scipy.md
 docs/quickstart.md
 examples/sparse_fock.py
@@ -35,6 +37,7 @@ scripts/run_simulator_tests.sh
 tests/conftest.py
 tests/test_bsr.py
 tests/test_formats.py
+tests/test_iterative.py
 tests/test_nki_bsr.py
 tests/test_nki_sim.py
 tests/test_nki_spmm.py
@@ -42,6 +45,7 @@ tests/test_ops.py
 tests/test_screening.py
 trnsparse/__init__.py
 trnsparse/formats.py
+trnsparse/iterative.py
 trnsparse/ops.py
 trnsparse/screening.py
 trnsparse.egg-info/PKG-INFO