trnsolver 0.3.0__tar.gz

trnsolver-0.3.0/.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,31 @@
+---
+name: Bug report
+about: Report a bug in trnsolver
+labels: bug
+---
+
+## Environment
+
+- trnsolver version:
+- Python version:
+- PyTorch version:
+- Hardware: CPU / trn1 / trn2 / inf2
+- OS:
+
+## Description
+
+What happened?
+
+## Steps to reproduce
+
+```python
+# Minimal code to reproduce
+```
+
+## Expected behavior
+
+What should have happened?
+
+## Actual behavior
+
+What happened instead? Include error messages/tracebacks if applicable.

trnsolver-0.3.0/.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,17 @@
+---
+name: Feature request
+about: Suggest a new feature or enhancement
+labels: enhancement
+---
+
+## Description
+
+What feature would you like?
+
+## Motivation
+
+Why is this needed? What problem does it solve?
+
+## Alternatives considered
+
+Any other approaches you've thought about?

trnsolver-0.3.0/.github/pull_request_template.md

@@ -0,0 +1,9 @@
+## Summary
+
+What does this PR do?
+
+## Test plan
+
+- [ ] Tests pass (`pytest tests/ -v`)
+- [ ] New tests added for new functionality
+- [ ] CHANGELOG.md updated (if user-facing change)

trnsolver-0.3.0/.github/workflows/ci.yml

@@ -0,0 +1,20 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+      - run: pip install -e ".[dev]"
+      - run: pytest tests/ -v -x --tb=short -m "not neuron"

trnsolver-0.3.0/.github/workflows/neuron.yml

@@ -0,0 +1,63 @@
+name: Neuron hardware tests
+
+# Manually triggered workflow that runs the `@pytest.mark.neuron`-marked suite
+# on a Trainium instance via SSM. The instance is provisioned separately with
+# `infra/terraform/`. This workflow wraps `scripts/run_neuron_tests.sh`.
+#
+# Prerequisites (one-time, manual):
+#   1. Provision the CI instance: `cd infra/terraform && terraform apply`.
+#   2. Register an AWS IAM role that trusts the GitHub OIDC provider for this
+#      repo. The role needs ec2:{Start,Stop,Describe}Instances and
+#      ssm:{SendCommand,GetCommandInvocation,DescribeInstanceInformation}.
+#   3. Add the role ARN to repo secrets as `AWS_ROLE_ARN`.
+#
+# Until step 3 is done, the workflow will fail at the OIDC assume-role step.
+
+on:
+  workflow_dispatch:
+    inputs:
+      instance_type:
+        description: "Instance family tag suffix (trn1, trn2, inf2)"
+        required: false
+        default: "trn1"
+      aws_region:
+        description: "AWS region"
+        required: false
+        default: "us-east-1"
+      warm:
+        description: "Run suite twice (warm NEFF cache pass)"
+        type: boolean
+        required: false
+        default: false
+
+concurrency:
+  # Only one hardware run at a time — the EC2 instance is shared.
+  group: neuron-hardware
+  cancel-in-progress: false
+
+jobs:
+  neuron:
+    runs-on: ubuntu-latest
+    timeout-minutes: 45
+    permissions:
+      id-token: write
+      contents: read
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Configure AWS credentials (OIDC)
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: ${{ secrets.AWS_ROLE_ARN }}
+          aws-region: ${{ inputs.aws_region }}
+
+      - name: Run neuron-marked tests via SSM
+        env:
+          AWS_REGION: ${{ inputs.aws_region }}
+        run: |
+          chmod +x scripts/run_neuron_tests.sh
+          if [[ "${{ inputs.warm }}" == "true" ]]; then
+            ./scripts/run_neuron_tests.sh --warm "${{ inputs.instance_type }}"
+          else
+            ./scripts/run_neuron_tests.sh "${{ inputs.instance_type }}"
+          fi

trnsolver-0.3.0/.github/workflows/notify-umbrella.yml

@@ -0,0 +1,36 @@
+name: Notify trnsci umbrella of docs change
+
+# Fires a repository_dispatch at trnsci/trnsci whenever this repo's docs or
+# mkdocs.yml change on main, so the combined site at trnsci.dev rebuilds
+# immediately rather than waiting for the daily cron.
+#
+# Requires repo/org secret TRNSCI_DISPATCH_TOKEN — a fine-grained PAT scoped
+# to the trnsci/trnsci repo with "Contents: Read and write" (which covers
+# repository_dispatch). See README for setup.
+
+on:
+  push:
+    branches: [main]
+    paths:
+      - "docs/**"
+      - "mkdocs.yml"
+      - "README.md"
+      - "CHANGELOG.md"
+  workflow_dispatch:
+
+jobs:
+  dispatch:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Trigger umbrella rebuild
+        env:
+          GH_TOKEN: ${{ secrets.TRNSCI_DISPATCH_TOKEN }}
+        run: |
+          if [ -z "${GH_TOKEN:-}" ]; then
+            echo "::notice::TRNSCI_DISPATCH_TOKEN not set; umbrella will pick up this change on next daily cron or manual dispatch."
+            exit 0
+          fi
+          gh api --method POST repos/trnsci/trnsci/dispatches \
+            -f event_type=subproject-updated \
+            -f "client_payload[repo]=${GITHUB_REPOSITORY}" \
+            -f "client_payload[sha]=${GITHUB_SHA}"

trnsolver-0.3.0/.github/workflows/publish.yml

@@ -0,0 +1,56 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - run: pip install build
+      - run: python -m build
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish-pypi:
+    needs: build
+    runs-on: ubuntu-latest
+    # Only publish stable releases (v1.2.3), not pre-releases (v1.2.3-rc1)
+    if: ${{ !contains(github.event.release.tag_name, '-') }}
+    environment:
+      name: pypi
+      url: https://pypi.org/project/trnsolver/
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1
+
+  publish-testpypi:
+    needs: build
+    runs-on: ubuntu-latest
+    # Only publish pre-releases to TestPyPI
+    if: ${{ contains(github.event.release.tag_name, '-') }}
+    environment:
+      name: testpypi
+      url: https://test.pypi.org/project/trnsolver/
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://test.pypi.org/legacy/

trnsolver-0.3.0/.gitignore

@@ -0,0 +1,26 @@
+__pycache__/
+*.py[cod]
+*$py.class
+*.egg-info/
+dist/
+build/
+.eggs/
+*.egg
+.pytest_cache/
+.tox/
+.venv/
+venv/
+*.so
+.idea/
+.vscode/
+*.swp
+*.swo
+*~
+site/
+.benchmarks/
+.terraform/
+*.tfstate
+*.tfstate.backup
+.terraform.lock.hcl
+.claude/
+.DS_Store

trnsolver-0.3.0/CHANGELOG.md

@@ -0,0 +1,87 @@
+# Changelog
+
+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).
+
+## [Unreleased]
+
+## [0.3.0] — 2026-04-12
+
+### Added
+
+- `inv_sqrt_spd_ns(A)` — Newton-Schulz iteration for `A^{-1/2}`. All-GEMM,
+  returns `(X, iters, residual)`. Progress on #14 (NKI-GEMM backend swap
+  lands when trnblas GEMM validates on hardware).
+- scipy.linalg baselines (eigh, cholesky, cho_solve) in
+  `benchmarks/bench_solver.py` for LAPACK comparison. Progress on #13.
+- `benchmarks/bench_cuda.py` — CUDA / cuSOLVER benchmarks, auto-skipped
+  when no GPU is available. Vintage-matched methodology (trn1 ↔ A10G,
+  trn2 ↔ H100). Progress on #13.
+- `infra/terraform/gpu.tf` — opt-in GPU CI instance (default `g5.xlarge`
+  A10G, vintage peer of trn1; `p5.4xlarge` for trn2 via override).
+  Disabled by default via `enable_gpu_ci = false`.
+- `scripts/run_cuda_tests.sh` — SSM-based runner mirroring
+  `run_neuron_tests.sh` for the GPU box.
+
+### Changed
+
+- `docs/benchmarks.md` rewritten with vintage-matching methodology and
+  real CPU numbers; `docs/aws_setup.md` documents the GPU instance and
+  H100 opt-in.
+- CI: bumped to `actions/checkout@v6` and `actions/setup-python@v6`
+  (Node.js 24 runtime); standalone Deploy Docs workflow removed — docs
+  are served from trnsci.dev.
+- `pyproject.toml` normalized across the trnsci suite; added
+  `Documentation` URL pointing to `trnsci.dev/trnsolver/`.
+
+## [0.2.0] — 2026-04-12
+
+### Added
+
+- `jacobi_preconditioner(A)` in `trnsolver.iterative` — diagonal preconditioner
+  callable for `cg`. Partial progress on #16.
+- Benchmark suite `benchmarks/bench_solver.py` covering eigh, factorizations,
+  solve, CG (with/without Jacobi preconditioner), and GMRES. Closes #8.
+- Status badges in README (CI, PyPI, Python versions, license, docs). Closes #4.
+- Manual Neuron hardware CI workflow (`.github/workflows/neuron.yml`) —
+  `workflow_dispatch`-only, OIDC-based, wraps `scripts/run_neuron_tests.sh`.
+  Closes #5.
+
+### Changed
+
+- `cg(A, b, M=...)` tensor preconditioner is now applied as `M @ r` (already-
+  inverted operator), not `inv(M) @ r` recomputed every iteration.
+- Repository transferred from `scttfrdmn/trnsolver` to `trnsci/trnsolver`.
+  Hardcoded URLs updated across docs, pyproject, terraform, and CI.
+
+## [0.1.1] — 2026-04-12
+
+### Changed
+
+- Bumped `neuronxcc` floor from `>=2.15` to `>=2.24` to unify with the
+  rest of the trnsci suite. `torch-neuronx` floor bumped to `>=2.9`.
+
+## [0.1.0] - 2026-04-11
+
+### Added
+
+- Initial scaffold of `trnsolver` — linear solvers and eigendecomposition
+  for AWS Trainium via NKI. Third sibling in the `trn-*` suite alongside
+  `trnblas` and `trnfft`.
+- Eigenvalue API: `eigh` (symmetric), `eigh_generalized` (Cholesky reduction
+  for `Ax = λBx`).
+- Factorizations: `cholesky`, `lu`, `qr`, `solve`, `solve_spd`, `inv_spd`,
+  `inv_sqrt_spd`.
+- Iterative solvers: `cg`, `gmres`.
+- NKI Jacobi rotation kernel scaffold in `trnsolver/nki/dispatch.py` with
+  auto/pytorch/nki backend selection. PyTorch fallback active until
+  on-hardware validation lands.
+- SCF demo (`examples/scf_eigen.py`) for the quantum-chemistry use case:
+  build Fock → solve `FC = SCε` → density → convergence loop.
+- pytest test suite covering eigen, factorization, and iterative solvers
+  on CPU.
+- mkdocs-material site, GitHub Actions CI/docs/publish workflows, Terraform
+  module for the Neuron CI instance, and `scripts/run_neuron_tests.sh`
+  SSM-driven hardware test runner — all mirroring the trnblas/trnfft layout.

trnsolver-0.3.0/CLAUDE.md

@@ -0,0 +1,105 @@
+# trnsolver
+
+Linear solvers and eigendecomposition for AWS Trainium via NKI.
+Part of the trnsci scientific computing suite.
+
+## What This Is
+
+Eigenvalue problems, matrix factorizations, and iterative solvers targeting
+Trainium's Tensor Engine. Complements trnblas (BLAS) and trnfft (FFT).
+
+**Primary use case:** SCF eigenvalue problem for quantum chemistry. The SCF
+iteration solves FC = SCε at each step — a generalized symmetric eigenvalue
+problem reduced to standard form via Cholesky of the overlap matrix S. This
+feeds into DF-MP2 (trnblas) for the correlation energy.
+
+## Architecture
+
+```
+trnsolver/
+├── trnsolver/
+│   ├── __init__.py          # Re-exports all solver operations
+│   ├── eigen.py             # eigh, eigh_generalized (Jacobi + torch fallback)
+│   ├── factor.py            # cholesky, lu, qr, solve, solve_spd, inv_sqrt_spd
+│   ├── iterative.py         # cg, gmres
+│   └── nki/
+│       ├── __init__.py
+│       └── dispatch.py      # auto/pytorch/nki dispatch + Jacobi rotation kernel
+├── tests/
+│   ├── conftest.py          # Fixtures: symmetric, SPD, random matrices
+│   ├── test_eigen.py        # Eigenvalue tests (reconstruction, orthogonality)
+│   ├── test_factor.py       # Cholesky, LU, QR, solve tests
+│   └── test_iterative.py    # CG, GMRES convergence tests
+├── examples/
+│   └── scf_eigen.py         # SCF iteration demo (quantum chemistry)
+├── pyproject.toml
+├── README.md
+├── LICENSE                  # Apache 2.0
+└── CLAUDE.md                # This file
+```
+
+## SCF → Solver Mapping
+
+| SCF Step | Math | trnsolver Call |
+|----------|------|---------------|
+| Overlap Cholesky | S = L L^T | `cholesky(S)` |
+| Reduce to standard | A' = L^{-1} F L^{-T} | `eigh_generalized(F, S)` |
+| Eigensolve | A' V = V Λ | `eigh(A')` (Jacobi on NKI) |
+| Back-transform | C = L^{-T} V | `trsm` (in trnblas) |
+| Density | P = C_occ @ C_occ^T | `syrk` (in trnblas) |
+| Metric inverse | J^{-1/2} | `inv_sqrt_spd(J)` |
+
+## NKI Jacobi Strategy
+
+The Jacobi eigensolver maps to Trainium hardware:
+
+- **Each Givens rotation** is a 2-row/col update — rank-2 matmul on Tensor Engine
+- **Off-diagonal max finding** maps to Vector Engine reduction
+- **Sweep convergence check** is a scalar reduction on Scalar Engine
+- **Eigenvector accumulation** V' = V @ G is a batched column update
+
+The Jacobi method is preferred over Householder tridiagonalization on
+Trainium because each rotation is a fixed-size operation that maps cleanly
+to the 128×128 tile, vs Householder which requires growing reflector chains.
+
+## Known Gaps & Design Notes
+
+- **NKI Jacobi kernel is a stub.** Falls back to torch.linalg.eigh until
+  validated on hardware. The rotation kernel scaffold is in nki/dispatch.py.
+
+- **Jacobi is O(n³) per sweep with O(n) sweeps** — cubic overall but with
+  a large constant. For n > ~500, Householder + QR iteration would be faster.
+  The Jacobi method is chosen because it's embarrassingly parallel (rotations
+  on disjoint pairs can run concurrently) and maps cleanly to NKI tiles.
+
+- **CG/GMRES are pure PyTorch.** The inner products and matvecs would
+  benefit from NKI Level 1/2 kernels in trnblas, but the iteration logic
+  stays on the host.
+
+- **inv_sqrt_spd uses eigendecomposition**, not the Newton-Schulz iteration.
+  Newton-Schulz (X_{k+1} = 0.5 X_k (3I - A X_k²)) is all GEMMs and would
+  map better to trnblas, but needs good initial scaling.
+
+## Dependencies
+
+- `torch>=2.1` — tensor operations and CPU fallback
+- `numpy>=1.24` — test reference
+- `neuronxcc` — NKI kernels (optional, only on Neuron hardware)
+
+## Development
+
+```bash
+pip install -e ".[dev]"
+pytest tests/ -v
+python examples/scf_eigen.py --demo
+```
+
+## Naming Convention
+
+Sibling repos in the trn-* suite:
+- `trnfft` — FFT + complex ops (https://github.com/trnsci/trnfft)
+- `trnblas` — BLAS operations (https://github.com/trnsci/trnblas)
+- `trnsolver` — Linear solvers, eigendecomposition (this repo)
+- `trnrand` — Random number generation (planned)
+
+All repos: Python/NKI, Apache 2.0.

trnsolver-0.3.0/CODE_OF_CONDUCT.md

@@ -0,0 +1,25 @@
+# Code of Conduct
+
+This project adopts the [Contributor Covenant](https://www.contributor-covenant.org/), version 2.1, as its code of conduct. The full text is available at:
+
+https://www.contributor-covenant.org/version/2/1/code_of_conduct/
+
+## Scope
+
+This Code applies within all project spaces — the GitHub repositories under the [trnsci](https://github.com/trnsci) organization, issue trackers, pull request discussions, and any public forums where an individual is representing the project or its community.
+
+## Reporting
+
+Concerns about behaviour in project spaces may be reported confidentially to:
+
+**scttfrdmn@gmail.com**
+
+All reports will be reviewed and investigated promptly and fairly. The project maintainers are obligated to respect the privacy and security of the reporter.
+
+## Enforcement
+
+Maintainers who do not follow or enforce the Code of Conduct in good faith may face temporary or permanent repercussions as determined by other members of the project's leadership.
+
+## Attribution
+
+This Code of Conduct is adapted from the Contributor Covenant, version 2.1, available at https://www.contributor-covenant.org/version/2/1/code_of_conduct.html. It is licensed under CC BY 4.0.

trnsolver-0.3.0/CONTRIBUTING.md

@@ -0,0 +1,69 @@
+# Contributing to trnsolver
+
+Thanks for your interest in contributing. This document covers the whole suite — the same guidelines apply to the umbrella and to each of the six sub-projects (`trnfft`, `trnblas`, `trnrand`, `trnsolver`, `trnsparse`, `trntensor`).
+
+## Where does my change belong?
+
+- **Scoped to a single library** — bug fix, new feature, added kernel in one sub-project → open a PR in that sub-project's repo.
+- **Cross-cutting** — integration examples, umbrella docs, CUDA-rosetta content, coordinated version bumps, shared tooling → open a PR against [trnsci/trnsci](https://github.com/trnsci/trnsci).
+
+If you aren't sure, open the PR where it's easiest and we'll move it.
+
+## Development setup
+
+```bash
+git clone git@github.com:trnsci/trnsci.git
+cd trnsci
+# Clone sibling sub-projects if you don't already have them:
+for p in trnfft trnblas trnrand trnsolver trnsparse trntensor; do
+  [ -d "$p" ] || git clone "git@github.com:trnsci/$p.git"
+done
+make install-dev
+make test-all
+```
+
+A single sub-project in isolation:
+
+```bash
+cd trnfft   # or any other
+pip install -e ".[dev]"
+pytest tests/ -v -m "not neuron"
+```
+
+## Tests
+
+- `pytest tests/ -v -m "not neuron"` — the CPU-only suite. Must pass on every PR.
+- `pytest tests/ -v -m neuron` — on-hardware suite, run from a machine with AWS credentials via `scripts/run_neuron_tests.sh`. Not part of standard CI.
+
+## Conventions
+
+- **License header** — new files get the Apache 2.0 notice and `Copyright 2026 Scott Friedman` where a notice is customary (long source files, not one-line configs).
+- **Python** — ≥ 3.10; `torch` ≥ 2.1; `numpy` ≥ 1.24.
+- **No emoji** in source, commits, or docs unless the feature is literally about emoji.
+- **Docstrings** only where they help a reader who doesn't have the context you had when writing. Don't restate the signature.
+- **No throwaway files** in the repo root — work-in-progress notes belong in PR descriptions or `docs/`.
+
+## Commit style
+
+Conventional-ish prefixes: `fix:`, `feat:`, `docs:`, `test:`, `bench:`, `chore:`, `refactor:`. Keep the subject line ≤ 72 characters. Longer context goes in the body.
+
+## Pull requests
+
+Before opening a PR:
+
+- [ ] `pytest tests/ -v -m "not neuron"` passes
+- [ ] If you changed public API or behaviour, `CHANGELOG.md` has an `[Unreleased]` entry
+- [ ] If you added a public symbol, the project's `docs/api.md` lists it
+- [ ] PR description says *what* changed and *why*, not just *what*
+
+## Issues
+
+Label each issue one of: `bug`, `feature`, `question`, `docs`, `infra`. Bugs include repro steps and expected-vs-actual output. Features describe the use case.
+
+## Code of conduct
+
+By participating, you agree to abide by the [Contributor Covenant v2.1](CODE_OF_CONDUCT.md). Report concerns to the contact listed there.
+
+## License
+
+Contributions are Apache 2.0, per the repository's [LICENSE](LICENSE). You attest you have the right to submit the code.