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

Install

pip install trnsolver

# With Neuron hardware support
pip install trnsolver[neuron]

Usage

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

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

Status

• Jacobi eigensolver with PyTorch backend
• Generalized eigenvalue problem (Cholesky reduction)
• Cholesky, LU, QR factorizations
• Direct and SPD solvers
• CG and GMRES iterative solvers
• 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

License

Apache 2.0 — Copyright 2026 Scott Friedman