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
