trnsolver 0.3.0__py3-none-any.whl

trnsolver/__init__.py

@@ -0,0 +1,33 @@
+"""
+trnsolver — Linear solvers and eigendecomposition for AWS Trainium via NKI.
+
+Eigenvalue problems, factorizations, and iterative solvers for scientific
+computing on Trainium. Complements trnblas (BLAS) and trnfft (FFT).
+Part of the trnsci scientific computing suite.
+"""
+
+__version__ = "0.3.0"
+
+# Eigenvalue decomposition
+from .eigen import eigh, eigh_generalized
+
+# Factorizations
+from .factor import cholesky, lu, qr, solve, solve_spd, inv_spd, inv_sqrt_spd, inv_sqrt_spd_ns
+
+# Iterative solvers
+from .iterative import cg, gmres, jacobi_preconditioner
+
+# Backend control
+from .nki import HAS_NKI, set_backend, get_backend
+
+__all__ = [
+    # Eigen
+    "eigh", "eigh_generalized",
+    # Factorization
+    "cholesky", "lu", "qr", "solve", "solve_spd", "inv_spd", "inv_sqrt_spd",
+    "inv_sqrt_spd_ns",
+    # Iterative
+    "cg", "gmres", "jacobi_preconditioner",
+    # Backend
+    "HAS_NKI", "set_backend", "get_backend",
+]

trnsolver/eigen.py

@@ -0,0 +1,170 @@
+"""
+Eigenvalue decomposition for Trainium.
+
+Symmetric eigenvalue problem: A @ V = V @ diag(eigenvalues)
+
+Methods:
+- Jacobi iteration (NKI-accelerable: rotation = 2×2 matmul per sweep)
+- torch.linalg.eigh fallback for non-NKI backends
+
+The Jacobi method is the natural fit for Trainium because each rotation
+is a small matmul that maps to the Tensor Engine, and the sweep over
+off-diagonal elements maps to the Vector Engine for max-finding.
+
+Primary use case: SCF eigenvalue problem FC = SCε in quantum chemistry.
+The generalized eigenproblem is first reduced to standard form via
+Cholesky: S = L L^T, then solve (L^{-1} F L^{-T}) C' = C' ε.
+"""
+
+from __future__ import annotations
+
+import math
+import torch
+from typing import Optional, Tuple
+
+from .nki import _use_nki
+
+
+def eigh(
+    A: torch.Tensor,
+    max_sweeps: int = 100,
+    tol: float = 1e-10,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Symmetric eigenvalue decomposition: A = V @ diag(w) @ V^T
+
+    Args:
+        A: Symmetric matrix (n, n)
+        max_sweeps: Maximum Jacobi sweeps
+        tol: Convergence threshold (sum of squared off-diagonal elements)
+
+    Returns:
+        eigenvalues: (n,) sorted ascending
+        eigenvectors: (n, n) columns are eigenvectors
+    """
+    if _use_nki():
+        return _jacobi_eigh(A, max_sweeps, tol)
+    return _torch_eigh(A)
+
+
+def eigh_generalized(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    max_sweeps: int = 100,
+    tol: float = 1e-10,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Generalized symmetric eigenvalue problem: A @ x = λ B @ x
+
+    Reduces to standard form via Cholesky of B:
+        B = L @ L^T
+        (L^{-1} A L^{-T}) y = λ y
+        x = L^{-T} y
+
+    This is the SCF path: solve FC = SCε where F is the Fock matrix
+    and S is the overlap matrix.
+    """
+    L = torch.linalg.cholesky(B)
+    # L^{-1} A L^{-T}
+    L_inv_A = torch.linalg.solve_triangular(L, A, upper=False)
+    A_prime = torch.linalg.solve_triangular(L, L_inv_A.T, upper=False).T
+
+    # Symmetrize (numerical stability)
+    A_prime = 0.5 * (A_prime + A_prime.T)
+
+    eigenvalues, V_prime = eigh(A_prime, max_sweeps, tol)
+
+    # Back-transform: V = L^{-T} @ V_prime
+    eigenvectors = torch.linalg.solve_triangular(L.T, V_prime, upper=True)
+
+    return eigenvalues, eigenvectors
+
+
+def _torch_eigh(A: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Fallback to torch.linalg.eigh."""
+    return torch.linalg.eigh(A)
+
+
+def _jacobi_eigh(
+    A: torch.Tensor,
+    max_sweeps: int = 100,
+    tol: float = 1e-10,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Classical Jacobi eigenvalue algorithm.
+
+    Each sweep applies n(n-1)/2 Givens rotations to zero off-diagonal elements.
+    Each rotation is a 2×2 orthogonal transform — maps to Tensor Engine on NKI.
+
+    Convergence: quadratic for well-separated eigenvalues.
+    """
+    n = A.shape[0]
+    assert A.shape == (n, n), f"Expected square matrix, got {A.shape}"
+
+    # Work on a copy
+    D = A.clone()
+    V = torch.eye(n, dtype=A.dtype, device=A.device)
+
+    for sweep in range(max_sweeps):
+        # Check convergence: sum of squared off-diagonal elements
+        off_diag = D.clone()
+        off_diag.fill_diagonal_(0.0)
+        off_norm = (off_diag ** 2).sum().item()
+        if off_norm < tol:
+            break
+
+        # Sweep over all off-diagonal pairs
+        for p in range(n):
+            for q in range(p + 1, n):
+                if abs(D[p, q].item()) < tol * 0.01:
+                    continue
+
+                # Compute rotation angle
+                if abs(D[p, p].item() - D[q, q].item()) < 1e-15:
+                    theta = math.pi / 4.0
+                else:
+                    tau = (D[q, q].item() - D[p, p].item()) / (2.0 * D[p, q].item())
+                    if tau >= 0:
+                        t = 1.0 / (tau + math.sqrt(1.0 + tau * tau))
+                    else:
+                        t = -1.0 / (-tau + math.sqrt(1.0 + tau * tau))
+                    c = 1.0 / math.sqrt(1.0 + t * t)
+                    s = t * c
+                    theta = math.atan2(s, c)
+
+                c = math.cos(theta)
+                s = math.sin(theta)
+
+                # Apply rotation to D: D' = G^T @ D @ G
+                # Only rows/cols p and q change
+                for r in range(n):
+                    if r == p or r == q:
+                        continue
+                    d_rp = D[r, p].item()
+                    d_rq = D[r, q].item()
+                    D[r, p] = c * d_rp - s * d_rq
+                    D[p, r] = D[r, p]
+                    D[r, q] = s * d_rp + c * d_rq
+                    D[q, r] = D[r, q]
+
+                d_pp = D[p, p].item()
+                d_qq = D[q, q].item()
+                d_pq = D[p, q].item()
+                D[p, p] = c * c * d_pp - 2 * s * c * d_pq + s * s * d_qq
+                D[q, q] = s * s * d_pp + 2 * s * c * d_pq + c * c * d_qq
+                D[p, q] = 0.0
+                D[q, p] = 0.0
+
+                # Accumulate eigenvectors: V' = V @ G
+                for r in range(n):
+                    v_rp = V[r, p].item()
+                    v_rq = V[r, q].item()
+                    V[r, p] = c * v_rp - s * v_rq
+                    V[r, q] = s * v_rp + c * v_rq
+
+    # Extract eigenvalues from diagonal
+    eigenvalues = torch.diag(D)
+
+    # Sort ascending
+    idx = torch.argsort(eigenvalues)
+    eigenvalues = eigenvalues[idx]
+    V = V[:, idx]
+
+    return eigenvalues, V

trnsolver/factor.py

@@ -0,0 +1,157 @@
+"""
+Matrix factorizations for Trainium.
+
+Cholesky, LU, QR — the building blocks for linear solvers.
+
+For DF-MP2:
+- Cholesky of the Coulomb metric J = L L^T
+- Triangular solve L X = B for density fitting coefficients
+- These feed into trnblas.trsm for the actual solve
+
+For SCF:
+- Cholesky of overlap matrix S for generalized eigenvalue reduction
+"""
+
+from __future__ import annotations
+
+import torch
+from typing import Tuple, Optional
+
+
+def cholesky(A: torch.Tensor, upper: bool = False) -> torch.Tensor:
+    """Cholesky factorization: A = L @ L^T (or A = U^T @ U if upper=True).
+
+    A must be symmetric positive definite.
+    """
+    L = torch.linalg.cholesky(A)
+    if upper:
+        return L.T
+    return L
+
+
+def lu(A: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """LU factorization with partial pivoting: P @ A = L @ U
+
+    Returns:
+        P: Permutation matrix (n, n)
+        L: Lower triangular with unit diagonal (n, n)
+        U: Upper triangular (n, n)
+    """
+    LU, pivots = torch.linalg.lu_factor(A)
+    P, L, U = torch.lu_unpack(LU, pivots)
+    return P, L, U
+
+
+def qr(A: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """QR factorization: A = Q @ R
+
+    Returns:
+        Q: Orthogonal matrix (m, min(m,n))
+        R: Upper triangular (min(m,n), n)
+    """
+    return torch.linalg.qr(A)
+
+
+def solve(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    """Solve A @ X = B for X.
+
+    Uses LU factorization internally.
+    """
+    return torch.linalg.solve(A, B)
+
+
+def solve_spd(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    """Solve A @ X = B where A is symmetric positive definite.
+
+    Uses Cholesky factorization (faster than LU for SPD).
+    """
+    L = cholesky(A)
+    squeeze = B.dim() == 1
+    if squeeze:
+        B = B.unsqueeze(1)
+    Y = torch.linalg.solve_triangular(L, B, upper=False)
+    X = torch.linalg.solve_triangular(L.T, Y, upper=True)
+    if squeeze:
+        X = X.squeeze(1)
+    return X
+
+
+def inv_spd(A: torch.Tensor) -> torch.Tensor:
+    """Inverse of symmetric positive definite matrix via Cholesky.
+
+    Used for J^{-1} in density fitting (though J^{-1/2} via solve is preferred).
+    """
+    n = A.shape[0]
+    return solve_spd(A, torch.eye(n, dtype=A.dtype, device=A.device))
+
+
+def inv_sqrt_spd(A: torch.Tensor) -> torch.Tensor:
+    """A^{-1/2} for symmetric positive definite A.
+
+    Computed via eigendecomposition: A = V diag(λ) V^T → A^{-1/2} = V diag(1/√λ) V^T
+
+    Used in density fitting for the metric contraction J^{-1/2}.
+    """
+    eigenvalues, V = torch.linalg.eigh(A)
+    # Clamp small eigenvalues for numerical stability
+    eigenvalues = torch.clamp(eigenvalues, min=1e-12)
+    inv_sqrt_eig = 1.0 / torch.sqrt(eigenvalues)
+    return V @ torch.diag(inv_sqrt_eig) @ V.T
+
+
+def inv_sqrt_spd_ns(
+    A: torch.Tensor,
+    *,
+    max_iters: int = 20,
+    tol: float = 1e-7,
+) -> Tuple[torch.Tensor, int, float]:
+    """A^{-1/2} for SPD A via the coupled Newton-Schulz iteration.
+
+    Iterates:
+        Y_{k+1} = Y_k T
+        Z_{k+1} = T Z_k
+    with T = 0.5 (3I - Z_k Y_k), starting from Y_0 = A / s, Z_0 = I where
+    s = ||A||_F is a scaling factor that keeps the spectrum of Y_0 inside
+    (0, 1) for convergence. At convergence Y_k → (A/s)^{1/2} and
+    Z_k → (A/s)^{-1/2}, so A^{-1/2} = Z_k / sqrt(s). Convergence is tracked
+    via ||Y_k Z_k - I||_F (which must go to 0).
+
+    All operations are GEMMs — the workload shape that maps well to the
+    Trainium Tensor Engine via trnblas. Prefer this over :func:`inv_sqrt_spd`
+    when the NKI backend is active and A is well-conditioned. Falls back
+    to the eigendecomposition path (:func:`inv_sqrt_spd`) is recommended
+    for A with eigenvalue spread beyond ~1e6.
+
+    Args:
+        A: Symmetric positive definite matrix (n, n).
+        max_iters: Maximum iterations.
+        tol: Relative Frobenius-norm tolerance on ||Y - I||.
+
+    Returns:
+        X: Approximation of A^{-1/2}.
+        iters: Number of iterations executed.
+        residual: Final ||Y - I||_F relative to ||I||_F.
+
+    Note:
+        TODO: power-iteration-based spectral-norm scaling is tighter than
+        Frobenius for matrices with one dominant eigenvalue. Frobenius is
+        good enough for typical DF-MP2 metric matrices.
+    """
+    n = A.shape[0]
+    s = torch.linalg.norm(A, ord="fro")
+    I = torch.eye(n, dtype=A.dtype, device=A.device)
+
+    Y = A / s
+    Z = I.clone()
+
+    norm_I = torch.linalg.norm(I, ord="fro")
+    residual = float("inf")
+    for k in range(max_iters):
+        T = 0.5 * (3.0 * I - Z @ Y)
+        Y = Y @ T
+        Z = T @ Z
+        residual = (torch.linalg.norm(Y @ Z - I, ord="fro") / norm_I).item()
+        if residual < tol:
+            return Z / torch.sqrt(s), k + 1, residual
+
+    return Z / torch.sqrt(s), max_iters, residual

trnsolver/iterative.py

@@ -0,0 +1,179 @@
+"""
+Iterative linear solvers for Trainium.
+
+Conjugate Gradient (CG) for SPD systems, GMRES for general systems.
+These are relevant for large-scale SCF with linear-scaling methods
+and for CPSCF (coupled-perturbed SCF) response equations.
+
+All inner products and matvecs map to trnblas Level 1/2 ops.
+"""
+
+from __future__ import annotations
+
+import torch
+from typing import Optional, Callable, Tuple
+
+
+def jacobi_preconditioner(A: torch.Tensor) -> Callable[[torch.Tensor], torch.Tensor]:
+    """Build a Jacobi (diagonal) preconditioner M(r) = r / diag(A).
+
+    Cheap and effective when A is diagonally dominant. Approximates A^{-1}
+    by inverting only the diagonal.
+    """
+    d = torch.diagonal(A).clone()
+    if torch.any(d.abs() < 1e-15):
+        raise ValueError("Jacobi preconditioner: diagonal has near-zero entries")
+    inv_d = 1.0 / d
+    return lambda r: r * inv_d
+
+
+def cg(
+    A: torch.Tensor | Callable,
+    b: torch.Tensor,
+    x0: Optional[torch.Tensor] = None,
+    tol: float = 1e-6,
+    maxiter: int = 1000,
+    M: Optional[torch.Tensor | Callable] = None,
+) -> Tuple[torch.Tensor, int, float]:
+    """Conjugate Gradient solver for A @ x = b (A must be SPD).
+
+    Args:
+        A: SPD matrix (n, n) or callable matvec A(x) → A@x
+        b: Right-hand side (n,)
+        x0: Initial guess (default: zeros)
+        tol: Convergence tolerance on relative residual
+        maxiter: Maximum iterations
+        M: Preconditioner as callable M(r) → approx A^{-1} @ r, or tensor
+            (interpreted as the already-inverted preconditioner applied as M@r).
+            Use `jacobi_preconditioner(A)` for the common diagonal case.
+
+    Returns:
+        x: Solution vector
+        iters: Number of iterations
+        residual: Final relative residual norm
+    """
+    n = b.shape[0]
+    matvec = A if callable(A) else lambda x: torch.mv(A, x)
+    if M is None:
+        precond = None
+    elif callable(M):
+        precond = M
+    else:
+        precond = lambda r: torch.mv(M, 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
+
+    if precond is not None:
+        z = precond(r)
+    else:
+        z = r.clone()
+
+    p = z.clone()
+    rz = torch.dot(r, z).item()
+
+    for k in range(maxiter):
+        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()
+        if r_norm / b_norm < tol:
+            return x, k + 1, r_norm / b_norm
+
+        if precond is not None:
+            z = precond(r)
+        else:
+            z = r.clone()
+
+        rz_new = torch.dot(r, z).item()
+        beta = rz_new / rz
+        p = z + beta * p
+        rz = rz_new
+
+    return x, maxiter, torch.linalg.norm(r).item() / b_norm
+
+
+def gmres(
+    A: torch.Tensor | Callable,
+    b: torch.Tensor,
+    x0: Optional[torch.Tensor] = None,
+    tol: float = 1e-6,
+    maxiter: int = 100,
+    restart: int = 30,
+) -> Tuple[torch.Tensor, int, float]:
+    """GMRES solver for A @ x = b (general, non-symmetric).
+
+    Restarted GMRES with Arnoldi process. Each iteration builds an
+    orthonormal Krylov basis via modified Gram-Schmidt, then solves
+    the least-squares Hessenberg system.
+
+    Args:
+        A: Matrix (n, n) or callable matvec
+        b: Right-hand side (n,)
+        x0: Initial guess
+        tol: Convergence tolerance
+        maxiter: Maximum outer iterations (restarts)
+        restart: Krylov subspace size before restart
+
+    Returns:
+        x: Solution vector
+        total_iters: Total matvec count
+        residual: Final relative residual norm
+    """
+    n = b.shape[0]
+    matvec = A if callable(A) else lambda x: torch.mv(A, x)
+
+    x = x0.clone() if x0 is not None else torch.zeros(n, dtype=b.dtype, device=b.device)
+    b_norm = torch.linalg.norm(b).item()
+    if b_norm < 1e-15:
+        return x, 0, 0.0
+
+    total_iters = 0
+
+    for outer in range(maxiter):
+        r = b - matvec(x)
+        r_norm = torch.linalg.norm(r).item()
+        if r_norm / b_norm < tol:
+            return x, total_iters, r_norm / b_norm
+
+        # Arnoldi process
+        m = min(restart, n)
+        V = torch.zeros(n, m + 1, dtype=b.dtype, device=b.device)
+        H = torch.zeros(m + 1, m, dtype=b.dtype, device=b.device)
+        V[:, 0] = r / r_norm
+        g = torch.zeros(m + 1, dtype=b.dtype, device=b.device)
+        g[0] = r_norm
+
+        for j in range(m):
+            total_iters += 1
+            w = matvec(V[:, j])
+
+            # Modified Gram-Schmidt
+            for i in range(j + 1):
+                H[i, j] = torch.dot(w, V[:, i])
+                w = w - H[i, j] * V[:, i]
+
+            H[j + 1, j] = torch.linalg.norm(w)
+            if H[j + 1, j].item() < 1e-14:
+                m = j + 1
+                break
+            V[:, j + 1] = w / H[j + 1, j]
+
+            # Check convergence via Givens rotations on g
+            # (simplified: just check residual norm)
+
+        # Solve least-squares: min ||H @ y - g||
+        y, _ = torch.linalg.lstsq(H[:m + 1, :m], g[:m + 1])[:2]
+        x = x + V[:, :m] @ y[:m]
+
+    r = b - matvec(x)
+    return x, total_iters, torch.linalg.norm(r).item() / b_norm

trnsolver/nki/__init__.py

@@ -0,0 +1,5 @@
+"""NKI kernel dispatch for Trainium solver acceleration."""
+
+from .dispatch import HAS_NKI, set_backend, get_backend, _use_nki
+
+__all__ = ["HAS_NKI", "set_backend", "get_backend", "_use_nki"]

trnsolver/nki/dispatch.py

@@ -0,0 +1,91 @@
+"""
+NKI dispatch for solver operations.
+
+The Jacobi eigensolver is the primary NKI acceleration target:
+each Givens rotation is a 2×2 matmul on the Tensor Engine,
+and the off-diagonal max-finding maps to the Vector Engine.
+"""
+
+from __future__ import annotations
+
+import os
+
+try:
+    import neuronxcc.nki as nki
+    import neuronxcc.nki.language as nl
+    import neuronxcc.nki.isa as nisa
+    HAS_NKI = True
+except ImportError:
+    HAS_NKI = False
+
+# When set, kernel-path failures re-raise instead of falling back to PyTorch.
+# Used by the validation suite to catch silent kernel breakage during iteration.
+_REQUIRE_NKI = os.environ.get("TRNSOLVER_REQUIRE_NKI", "").lower() in ("1", "true", "yes")
+
+_backend = "auto"
+
+
+def set_backend(backend: str):
+    global _backend
+    assert backend in ("auto", "pytorch", "nki")
+    if backend == "nki" and not HAS_NKI:
+        raise RuntimeError("NKI backend requires neuronxcc")
+    _backend = backend
+
+
+def get_backend() -> str:
+    return _backend
+
+
+def _use_nki() -> bool:
+    if _backend == "nki":
+        return True
+    if _backend == "pytorch":
+        return False
+    return HAS_NKI
+
+
+if HAS_NKI:
+
+    @nki.jit
+    def jacobi_rotation_kernel(
+        D_ref, V_ref, p: int, q: int, c: float, s: float, n: int
+    ):
+        """Apply single Jacobi rotation to D and V matrices.
+
+        D' = G^T @ D @ G  (updates rows/cols p, q of D)
+        V' = V @ G        (accumulates eigenvectors)
+
+        G is the Givens rotation matrix with cos(θ)=c, sin(θ)=s
+        at positions (p,q).
+
+        On Tensor Engine: the 2-row/col update is a rank-2 operation.
+        On Vector Engine: the element-wise updates along rows p, q.
+
+        STUB: Scaffolded for on-hardware validation.
+        """
+        # Load rows p and q of D
+        d_p = nl.load(D_ref[p, :])
+        d_q = nl.load(D_ref[q, :])
+
+        # Rotate rows
+        new_p = nl.add(nl.multiply(d_p, c), nl.multiply(d_q, -s))
+        new_q = nl.add(nl.multiply(d_p, s), nl.multiply(d_q, c))
+        nl.store(D_ref[p, :], new_p)
+        nl.store(D_ref[q, :], new_q)
+
+        # Rotate columns (D is symmetric, so update cols too)
+        d_col_p = nl.load(D_ref[:, p])
+        d_col_q = nl.load(D_ref[:, q])
+        nl.store(D_ref[:, p], nl.add(nl.multiply(d_col_p, c), nl.multiply(d_col_q, -s)))
+        nl.store(D_ref[:, q], nl.add(nl.multiply(d_col_p, s), nl.multiply(d_col_q, c)))
+
+        # Zero the (p,q) and (q,p) entries
+        nl.store(D_ref[p, q], 0.0)
+        nl.store(D_ref[q, p], 0.0)
+
+        # Accumulate eigenvectors: V' = V @ G
+        v_p = nl.load(V_ref[:, p])
+        v_q = nl.load(V_ref[:, q])
+        nl.store(V_ref[:, p], nl.add(nl.multiply(v_p, c), nl.multiply(v_q, -s)))
+        nl.store(V_ref[:, q], nl.add(nl.multiply(v_p, s), nl.multiply(v_q, c)))

trnsolver-0.3.0.dist-info/METADATA

@@ -0,0 +1,127 @@
+Metadata-Version: 2.4
+Name: trnsolver
+Version: 0.3.0
+Summary: Linear solvers and eigendecomposition for AWS Trainium via NKI
+Author-email: Scott Friedman <scttfrdmn@gmail.com>
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/trnsci/trnsolver
+Project-URL: Documentation, https://trnsci.dev/trnsolver/
+Project-URL: Repository, https://github.com/trnsci/trnsolver
+Project-URL: Issues, https://github.com/trnsci/trnsolver/issues
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.1
+Requires-Dist: numpy>=1.24
+Provides-Extra: neuron
+Requires-Dist: neuronxcc>=2.24; extra == "neuron"
+Requires-Dist: torch-neuronx>=2.9; extra == "neuron"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-benchmark>=4.0; extra == "dev"
+Requires-Dist: scipy>=1.11; extra == "dev"
+Dynamic: license-file
+
+# trnsolver
+
+[![CI](https://github.com/trnsci/trnsolver/actions/workflows/ci.yml/badge.svg)](https://github.com/trnsci/trnsolver/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/trnsolver)](https://pypi.org/project/trnsolver/)
+[![Python](https://img.shields.io/pypi/pyversions/trnsolver)](https://pypi.org/project/trnsolver/)
+[![License](https://img.shields.io/github/license/trnsci/trnsolver)](LICENSE)
+[![Docs](https://img.shields.io/badge/docs-mkdocs-blue)](https://trnsci.github.io/trnsolver/)
+
+Linear solvers and eigendecomposition for AWS Trainium via NKI.
+
+Eigenvalue problems, matrix factorizations, and iterative solvers for scientific computing on Trainium. The Jacobi eigensolver is the primary NKI acceleration target — each Givens rotation maps to a 2-row matmul on the Tensor Engine.
+
+Part of the trnsci scientific computing suite ([github.com/trnsci](https://github.com/trnsci)).
+
+## Install
+
+```bash
+pip install trnsolver
+
+# With Neuron hardware support
+pip install trnsolver[neuron]
+```
+
+## Usage
+
+```python
+import torch
+import trnsolver
+
+# Symmetric eigenvalue decomposition (Jacobi on NKI)
+eigenvalues, eigenvectors = trnsolver.eigh(A)
+
+# Generalized eigenproblem: A x = λ B x (the SCF problem)
+eigenvalues, eigenvectors = trnsolver.eigh_generalized(F, S)
+
+# Factorizations
+L = trnsolver.cholesky(A)
+P, L, U = trnsolver.lu(A)
+Q, R = trnsolver.qr(A)
+
+# Direct solvers
+x = trnsolver.solve(A, b)
+x = trnsolver.solve_spd(A, b)        # Cholesky-based (faster for SPD)
+A_inv_sqrt = trnsolver.inv_sqrt_spd(A)  # A^{-1/2} for density fitting
+
+# Iterative solvers
+x, iters, residual = trnsolver.cg(A, b, tol=1e-8)
+x, iters, residual = trnsolver.gmres(A, b, tol=1e-6)
+```
+
+## SCF Example
+
+```bash
+python examples/scf_eigen.py --demo
+python examples/scf_eigen.py --nbasis 50 --nocc 10
+```
+
+Demonstrates the self-consistent field iteration: build Fock matrix → solve generalized eigenproblem FC = SCε → build density → check convergence.
+
+## Operations
+
+| Category | Operation | Description |
+|----------|-----------|-------------|
+| Eigen | `eigh` | Symmetric eigendecomposition (Jacobi / torch) |
+| Eigen | `eigh_generalized` | Generalized: Ax = λBx via Cholesky reduction |
+| Factor | `cholesky` | A = LL^T |
+| Factor | `lu` | PA = LU |
+| Factor | `qr` | A = QR |
+| Solve | `solve` | Ax = b (LU-based) |
+| Solve | `solve_spd` | Ax = b (Cholesky, A is SPD) |
+| Solve | `inv_spd` | A^{-1} for SPD A |
+| Solve | `inv_sqrt_spd` | A^{-1/2} for density fitting metric |
+| Iterative | `cg` | Conjugate Gradient (SPD systems) |
+| Iterative | `gmres` | GMRES (general systems) |
+
+## Status
+
+- [x] Jacobi eigensolver with PyTorch backend
+- [x] Generalized eigenvalue problem (Cholesky reduction)
+- [x] Cholesky, LU, QR factorizations
+- [x] Direct and SPD solvers
+- [x] CG and GMRES iterative solvers
+- [x] SCF example (quantum chemistry use case)
+- [ ] NKI Jacobi rotation kernel validation
+- [ ] Newton-Schulz A^{-1/2} via trnblas GEMM
+- [ ] Benchmarks vs LAPACK/cuSOLVER
+
+## Related Projects
+
+| Project | What |
+|---------|------|
+| [trnfft](https://github.com/trnsci/trnfft) | FFT + complex ops |
+| [trnblas](https://github.com/trnsci/trnblas) | BLAS operations |
+| trnsolver | This repo |
+
+## License
+
+Apache 2.0 — Copyright 2026 Scott Friedman

trnsolver-0.3.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+trnsolver/__init__.py,sha256=_b9l0h05EWXkfBx6-vqONlJXEczM4LqmBe5Nwv6BIkM,927
+trnsolver/eigen.py,sha256=DNubX9zGe2_oY5OirtgZ_MkyYx6usLs01NtJ8sHaNlQ,5414
+trnsolver/factor.py,sha256=ocUPIK8vkPdPpa_gmbxTq5DAoY3HrM6b5B02LUEhhtg,4808
+trnsolver/iterative.py,sha256=J3NHCirwY0YhD2GnrqQlguCzRH-GgSI5NU4CP5obCvE,5333
+trnsolver/nki/__init__.py,sha256=931tvEv0-8H-9urnDd9KkSmAUdlxHiCKeBvuDXFkLek,192
+trnsolver/nki/dispatch.py,sha256=aJowSyPncJWWDA-A3tmF7U96BGQ5aSvBdVWMBO2OBOw,2780
+trnsolver-0.3.0.dist-info/licenses/LICENSE,sha256=hTt_varGm4HTIVm8GQ1pGR7jah4vW5IrddFIu94FPAE,11175
+trnsolver-0.3.0.dist-info/METADATA,sha256=vLVlAT40npUx8DEAhPGhKB_2crSX-GQs73P7--ziclM,4420
+trnsolver-0.3.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+trnsolver-0.3.0.dist-info/top_level.txt,sha256=GSWEYFB46Xyi7LrVpr5PKwtA_huGp8HY26_E-668_N4,10
+trnsolver-0.3.0.dist-info/RECORD,,

trnsolver-0.3.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

trnsolver-0.3.0.dist-info/licenses/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   Copyright 2026 Scott Friedman
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for describing the origin of the Work and
+      reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may accept support,
+      obligations, or other liability only on behalf of Yourself and on
+      Your sole responsibility, not on behalf of any other Contributor,
+      and only if You agree to indemnify, defend, and hold each Contributor
+      harmless for any liability incurred by, or claims asserted against,
+      such Contributor by reason of your accepting any such warranty or
+      additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2026 Scott Friedman
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

trnsolver-0.3.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+trnsolver