trntensor 0.1.2__tar.gz → 0.3.0__tar.gz

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

@@ -0,0 +1,59 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+# Opt into Node.js 24 for all third-party JavaScript actions ahead of
+# GitHub's June 2026 default switch. Silences the Node 20 deprecation
+# notice emitted by setup-uv and codecov-action's github-script dep.
+env:
+  FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v6
+      - run: uvx ruff check .
+      - run: uvx ruff format --check .
+
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+      - run: uv sync --extra dev
+      - run: uv run pytest tests/ -v -x --tb=short -m "not neuron and not nki_simulator" --cov=trntensor --cov-report=xml
+
+      - name: Upload coverage reports to Codecov
+        uses: codecov/codecov-action@v5
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          slug: trnsci/trntensor
+
+  nki-simulator:
+    # Runs NKI kernels through nki.simulate(kernel)(numpy_args) on CPU.
+    # Catches Python-trace-level errors (bad kwargs, dropped ops, shape
+    # mismatches) pre-merge without AWS round-trips. MLIR verifier
+    # errors remain hardware-only (simulator explicitly skips compile).
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v6
+        with:
+          python-version: "3.12"
+      - run: uv sync --extra dev
+      - name: Install NKI simulator
+        run: uv pip install --extra-index-url https://pip.repos.neuron.amazonaws.com "nki>=0.3.0"
+      - name: Run simulator-backed kernel tests
+        env:
+          TRNTENSOR_USE_SIMULATOR: "1"
+        run: uv run pytest tests/ -v -m nki_simulator --tb=short

{trntensor-0.1.2 → trntensor-0.3.0}/.github/workflows/notify-umbrella.yml

@@ -18,6 +18,9 @@ on:
       - "CHANGELOG.md"
   workflow_dispatch:
 
+env:
+  FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true
+
 jobs:
   dispatch:
     runs-on: ubuntu-latest

{trntensor-0.1.2 → trntensor-0.3.0}/.github/workflows/publish.yml

@@ -4,6 +4,9 @@ on:
   release:
     types: [published]
 
+env:
+  FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true
+
 jobs:
   build:
     runs-on: ubuntu-latest

{trntensor-0.1.2 → trntensor-0.3.0}/.gitignore

@@ -31,3 +31,9 @@ site/
 # coverage artifacts
 .coverage
 htmlcov/
+
+# terraform local state
+infra/terraform/.terraform/
+infra/terraform/terraform.tfstate*
+infra/terraform/.terraform.lock.hcl
+benchmark_results.json

trntensor-0.3.0/.pre-commit-config.yaml

@@ -0,0 +1,11 @@
+# Install locally: pip install pre-commit && pre-commit install
+# Runs on `git commit`; ruff auto-fixes staged Python files then formats them.
+# CI re-checks the same rules in the lint job — pre-commit keeps CI green.
+
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.6.9
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format

trntensor-0.3.0/CHANGELOG.md

@@ -0,0 +1,201 @@
+# 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]
+
+### Added
+
+- **`dtype` mixed-precision override for `einsum`** (#22) — accepts string
+  aliases (`"bf16"`, `"fp16"`, `"f32"`) or `torch.dtype` instances. When set,
+  all operands are cast to the requested dtype before contracting and the
+  result is returned in that dtype. Matches Neuron SDK autocast recommendations;
+  use `dtype="bf16"` to route fp32 models through the NKI bf16 matmul path
+  without changing the model's weight dtype.
+
+## [0.3.0] — 2026-04-16
+
+### Added
+
+- **Greedy contraction-path search for 3+ operand `einsum`** (#18) —
+  `plan_contraction()` now returns `strategy="path"` for three or more
+  operands. `_greedy_path_search` selects the cheapest binary contraction
+  order by minimizing per-step FLOP cost; the resulting
+  `ContractionPlan.contraction_path` drives `_execute_path`, which
+  routes each binary step through the full backend selection stack so
+  large sub-contractions still dispatch to NKI.
+
+- **`multi_einsum` shared-operand XLA residency** (#19) — When NKI
+  dispatch is active, `multi_einsum` detects operand tensors that
+  appear in more than one contraction (by object identity) and
+  pre-pins them to the XLA device once before executing the loop.
+  Eliminates redundant host↔device transfers for workloads like DF-MP2
+  where the three-center integral tensor `B` feeds many pair
+  contractions. Falls back to the existing per-contraction loop on CPU.
+
+- **Subscript and shape validation with descriptive errors** (#26) —
+  `plan_contraction()` now validates subscripts up-front and raises
+  `ValueError` with precise messages: wrong operand count, rank
+  mismatch, and inconsistent index sizes are all caught before any
+  torch operation runs. Eliminates cryptic downstream errors from
+  `torch.einsum`.
+
+- **PEP 561 `py.typed` marker** (#25) — `trntensor` now ships a
+  `py.typed` file so type checkers (mypy, pyright, etc.) recognise
+  the package as typed and apply inline annotations.
+
+- **`alpha`/`beta` scaling for `einsum`** (#20) — matches cuTENSOR's
+  GEMM-style interface:
+  `einsum(subscripts, A, B, alpha=α, beta=β, out=C)` returns
+  `α * contract(A, B) + β * C`. Defaults (`alpha=1, beta=0, out=None`)
+  preserve existing behaviour exactly. Useful for accumulation patterns
+  and in-place gradient updates without an extra allocation.
+
+- **Contraction plan cache** (#29 partial) — `plan_contraction()` caches
+  results by `(subscripts, operand shapes)`. Repeated calls with the same
+  subscript and shapes skip replanning entirely. `clear_plan_cache()` and
+  `plan_cache_info()` are exported from the top-level API.
+
+- **Tensor Train (TT) decomposition** (#23) — `tt_decompose(tensor, max_rank)`
+  decomposes a d-dimensional tensor into a chain of 3-tensor cores via TT-SVD
+  (Oseledets 2011), bond dimension capped at `max_rank`. `tt_reconstruct(cores)`
+  contracts the chain back. Useful for DMRG-style high-dimensional compression.
+
+- **Non-negative CP and warm-start CP** (#24) — `cp_decompose` gains two new
+  keyword arguments: `nonneg=True` switches ALS to multiplicative updates
+  to enforce non-negative factors; `factors=` accepts a list of pre-computed
+  factor matrices to warm-start from, skipping random initialization.
+  Both options compose.
+
+## [0.2.0] — 2026-04-15
+
+### Added
+
+- **`trntensor.to_xla(tensor)` / `trntensor.from_xla(tensor)`** — explicit
+  operand residency on the Trainium XLA device. Pre-pinning operands
+  lets repeated trntensor calls skip per-dispatch host↔device
+  transfer, which otherwise dominates at current kernel sizes. The
+  full DF-MP2 pipeline (`ao_to_mo_transform` → `mp2_energy`) with all
+  operands pre-pinned pays transfer cost once instead of once per
+  call. The dispatch layer's `_to_xla` helper takes a fast path when
+  every operand is already on XLA, returning the result on XLA — the
+  caller decides when to pull back via `from_xla`. Closes #34.
+- **`trntensor.ao_to_mo_transform(eri, C_occ, C_vir)`** — fused 4-index
+  AO→MO integral transform with K-tiling over the basis index (#37).
+  One NKI program computes `B[i,a,P] = Σ_{μν} C_occ[μ,i] · C_vir[ν,a] ·
+  eri[μ,ν,P]`. Tiles over μ (step 1) and ν (step 2) in TILE_K=128
+  chunks so `nbasis` up to 512 is supported; dispatch pads to the
+  nearest TILE_K multiple. Shape constraints: `nbasis ≤ 512`,
+  `nocc ≤ 128`, `nvir ≤ 512`. Composes with `mp2_energy` for the full
+  DF-MP2 pipeline from AO integrals to correlation energy. Validated
+  on trn1 (hardware) and via the CPU simulator CI job.
+- **NKI CPU simulator dispatch** via `TRNTENSOR_USE_SIMULATOR=1`.
+  Routes kernels through `nki.simulate(kernel)(numpy_args)` on CPU,
+  bypassing `torch_xla` + NEFF compile. Iteration loop drops from
+  ~5 min per SSM round-trip to seconds. Correctness-only — MLIR
+  verifier errors remain hardware-only.
+- **`nki-simulator` CI job on `ubuntu-latest`** — runs the
+  `nki_simulator`-marked suite against `nki>=0.3.0` from the AWS pip
+  index on every push + PR. Zero AWS cost for the correctness gate.
+- `tests/test_nki_sim.py` — simulator-backed correctness suite,
+  marker `nki_simulator`. Covers matmul, batched matmul,
+  `ao_to_mo_transform` (including K-tiled nbasis=256 and non-aligned
+  nbasis=200), and `mp2_energy`.
+- `scripts/run_simulator_tests.sh` — SSM runner for the simulator
+  suite on the trn1 DLAMI.
+- `docs/developing_kernels.md` — NKI kernel development guide with
+  trntensor-specific env vars and file locations.
+
+### Changed
+
+- **Migrated to NKI 0.3.0 / Neuron SDK 2.29.** Canonical `nki.*`
+  namespace; the legacy `neuronxcc.nki.*` shim is no longer used.
+  Kernels updated for the NKI 0.3.0 breaking-change surface:
+  `nisa.nc_matmul(dst=, stationary=, moving=, accumulate=True)` (all
+  kwargs); `nl.copy(psum)` returns a view — use `nl.ndarray` +
+  `nisa.tensor_copy` instead; tensor-tensor `nl.divide` dropped — use
+  `multiply × reciprocal`.
+- Dev workflow migrated to [uv](https://github.com/astral-sh/uv).
+  `uv sync --extra dev` replaces `pip install -e ".[dev]"`; CI uses
+  `astral-sh/setup-uv@v6` and `uv run pytest` / `uvx ruff`. `uv.lock`
+  is committed for reproducible installs.
+- Removed the `[neuron]` optional-dependencies extra. `nki` is
+  installed from the AWS Neuron pip index in CI or provided by the
+  Deep Learning AMI's pre-built venv on hardware.
+- `CONTRIBUTING.md` updated to reflect the uv-based setup.
+
+### Fixed
+
+- `mp2_energy_kernel` 1D-load ambiguity on pre-pinned XLA ε inputs
+  (#38). Reshape `eps_occ` / `eps_vir` to 2D `(N, 1)` at the dispatch
+  boundary; partition-dim inference is unambiguous regardless of
+  residency state.
+- `mp2_energy_kernel` 0-D SBUF rejection (`SBUF tensors must have at
+  least 2 dimensions`). Per-`(i,j)` reduction now uses a persistent
+  `(1, 1)` SBUF accumulator (`nl.zeros((1, 1), ...)`) instead of a
+  direct `nl.sum` store.
+- `_to_xla` fast-path now calls `xm.mark_step()` when operands are
+  already on XLA, forcing pending lazy computations to materialize
+  before the next kernel dispatch.
+
+### Known limitation
+
+- The full DF-MP2 pipeline (`ao_to_mo_transform` → `mp2_energy`) with
+  every operand pre-pinned exposes an NKI compiler bug on trn1: the
+  combined XLA lazy graph provokes `trn2-only shared memory`
+  instructions that fail verification on trn1. Workaround: `from_xla`
+  the intermediate `B` between the two calls. Tracked in #39 for
+  upstream AWS escalation.
+
+## [0.1.2] — 2026-04-12
+
+### Changed
+
+- `set_backend("nki")` now raises `RuntimeError` on non-Neuron hosts
+  instead of silently accepting the backend and failing later. Matches
+  the sibling-suite pattern.
+- CI actions bumped to `actions/checkout@v6` + `actions/setup-python@v6`
+  (Node.js 24), ahead of GitHub's June 2026 default switch.
+- pyproject metadata normalized across the trnsci suite (author email,
+  URLs, classifier list).
+- Standalone `docs.yml` removed — docs are now served via `trnsci.dev`
+  through the umbrella's combined build. `notify-umbrella.yml` pings
+  the umbrella on docs changes.
+- `infra/terraform/main.tf`: user-data clone URL corrected to
+  `trnsci/trntensor`.
+
+### Added
+
+- `benchmarks/bench_einsum.py` — pytest-benchmark cases for einsum
+  dispatch and decompositions. CPU baseline numbers populated in
+  `docs/benchmarks.md`.
+- `tests/test_nki.py` — backend-dispatch unit tests (CPU path).
+
+## [0.1.1] — 2026-04-12
+
+### Added
+
+- mkdocs site with `index`, `installation`, `quickstart`, `api`, `architecture`, `aws_setup`
+- `infra/terraform/` for on-hardware CI instance provisioning
+- `scripts/run_neuron_tests.sh` and benchmark helpers
+- GitHub Actions `ci.yml`, `docs.yml`, `publish.yml`
+- `Issues` and `Documentation` URLs in pyproject.toml
+- `tests/test_plan.py` — dedicated planner unit tests (parsing,
+  strategy selection, FLOP estimates); extended CP / Tucker coverage
+  (all-zero tensor, rank > min dim, unequal mode ranks)
+
+### 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-12
+
+### Added
+
+- Initial scaffold: einsum with contraction planning, CP / Tucker decompositions
+- NKI dispatch with fused-contraction kernel stubs
+- `examples/df_mp2_einsum.py` — DF-MP2 energy via einsum

{trntensor-0.1.2 → trntensor-0.3.0}/CLAUDE.md

@@ -59,3 +59,15 @@ pip install -e ".[dev]"
 pytest tests/ -v
 python examples/df_mp2_einsum.py --demo
 ```
+
+## Blog posts
+
+When asked to draft a blog post for this library for the [trnsci blog](https://trnsci.dev/blog/):
+
+1. Read the editorial brief at [`docs/blog/AUTHOR_BRIEF.md`](https://github.com/trnsci/trnsci/blob/main/docs/blog/AUTHOR_BRIEF.md) in the umbrella repo (trnsci/trnsci). It defines voice (authorless, library-as-subject), stance (architecture-first, transparency-always), and the nine required section headings.
+
+2. Find the prompt block for this library in [`BLOG_PROMPTS.md`](https://github.com/trnsci/trnsci/blob/main/BLOG_PROMPTS.md) at the umbrella repo root. It carries library-specific context and suggested architectural angles.
+
+3. Draft the post following the brief. Open a PR against `trnsci/trnsci` at `docs/blog/posts/<YYYY-MM-DD>-<slug>.md`. Scott (suite director) reviews before merge.
+
+The umbrella repo — not this one — owns the blog. Per-library retrospective posts are unsigned; library is the subject, no byline. See the brief for the full set of rules.

{trntensor-0.1.2 → trntensor-0.3.0}/CONTRIBUTING.md

@@ -26,8 +26,9 @@ A single sub-project in isolation:
 
 ```bash
 cd trnfft   # or any other
-pip install -e ".[dev]"
-pytest tests/ -v -m "not neuron"
+uv sync --extra dev
+uv run pre-commit install   # ruff check + format on every commit
+uv run pytest tests/ -v -m "not neuron"
 ```
 
 ## Tests

trntensor-0.3.0/PKG-INFO

@@ -0,0 +1,143 @@
+Metadata-Version: 2.4
+Name: trntensor
+Version: 0.3.0
+Summary: Tensor contractions for AWS Trainium via NKI
+Author-email: Scott Friedman <scttfrdmn@gmail.com>
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/trnsci/trntensor
+Project-URL: Documentation, https://trnsci.dev/trntensor/
+Project-URL: Repository, https://github.com/trnsci/trntensor
+Project-URL: Issues, https://github.com/trnsci/trntensor/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: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-benchmark>=4.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0; extra == "dev"
+Requires-Dist: ruff>=0.6; extra == "dev"
+Requires-Dist: pre-commit>=3.7; extra == "dev"
+Requires-Dist: scipy>=1.11; extra == "dev"
+Requires-Dist: mkdocs>=1.5; extra == "dev"
+Requires-Dist: mkdocs-material>=9.0; extra == "dev"
+Dynamic: license-file
+
+# trntensor
+
+[![CI](https://github.com/trnsci/trntensor/actions/workflows/ci.yml/badge.svg)](https://github.com/trnsci/trntensor/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/trnsci/trntensor/graph/badge.svg)](https://codecov.io/gh/trnsci/trntensor)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+[![PyPI](https://img.shields.io/pypi/v/trntensor)](https://pypi.org/project/trntensor/)
+[![Python](https://img.shields.io/pypi/pyversions/trntensor)](https://pypi.org/project/trntensor/)
+[![License](https://img.shields.io/github/license/trnsci/trntensor)](LICENSE)
+[![Docs](https://img.shields.io/badge/docs-trnsci.dev-blue)](https://trnsci.dev/trntensor/)
+
+Tensor contractions for AWS Trainium via NKI.
+
+Einstein summation with contraction planning, CP and Tucker decompositions. Expresses scientific tensor workloads naturally instead of decomposing to GEMM. Part of the trnsci scientific computing suite ([github.com/trnsci](https://github.com/trnsci)).
+
+## Current phase
+
+trntensor follows the [trnsci 5-phase roadmap](https://trnsci.dev/roadmap/). Active work is tracked in phase-labeled GitHub issues:
+
+- **[Phase 1 — correctness](https://github.com/trnsci/trntensor/issues/27)** (active): matmul + batched-matmul NKI kernels in place; awaiting hardware validation + additional `@pytest.mark.neuron` coverage.
+- **[Phase 2 — precision](https://github.com/trnsci/trntensor/issues/28)**: precision-aware contraction path selection (depends on [trnblas#22](https://github.com/trnsci/trnblas/issues/22) double-double GEMM).
+- **[Phase 3 — perf](https://github.com/trnsci/trntensor/issues/29)**: opt_einsum-style path planner, plan cache reuse.
+- **[Phase 4 — multi-chip](https://github.com/trnsci/trntensor/issues/30)**: sharded tensor contractions.
+- **[Phase 5 — generation](https://github.com/trnsci/trntensor/issues/31)**: trn2 fused multi-contraction paths.
+
+Suite-wide tracker: [trnsci/trnsci#1](https://github.com/trnsci/trnsci/issues/1).
+
+## Install
+
+```bash
+pip install trntensor
+```
+
+## Usage
+
+```python
+import torch
+import trntensor
+
+# Einsum — drop-in for torch.einsum with contraction planning
+C = trntensor.einsum("ij,jk->ik", A, B)           # matmul
+T = trntensor.einsum("ap,bp->ab", B_i, B_j)       # DF-MP2 contraction
+X = trntensor.einsum("mi,mnP->inP", C_occ, eri)   # AO→MO transform
+
+# Contraction planning
+plan = trntensor.plan_contraction("ij,jk->ik", A, B)
+flops = trntensor.estimate_flops("ij,jk->ik", A, B)
+
+# CP decomposition (tensor hypercontraction)
+factors, weights = trntensor.cp_decompose(tensor, rank=10)
+reconstructed = trntensor.cp_reconstruct(factors, weights)
+
+# Tucker decomposition (HOSVD)
+core, factors = trntensor.tucker_decompose(tensor, ranks=(5, 5, 5))
+```
+
+## Operations
+
+| Category | Operation | Description |
+|----------|-----------|-------------|
+| Contraction | `einsum` | General tensor contraction |
+| Contraction | `multi_einsum` | Multiple contractions (fusion-ready) |
+| Planning | `plan_contraction` | Analyze and select strategy |
+| Planning | `estimate_flops` | FLOPs for a contraction |
+| Decomposition | `cp_decompose` | CP/PARAFAC via ALS |
+| Decomposition | `tucker_decompose` | Tucker via HOSVD |
+
+## Status
+
+- [x] Einsum with matmul/bmm/torch dispatch
+- [x] Contraction planner
+- [x] CP decomposition (ALS)
+- [x] Tucker decomposition (HOSVD)
+- [x] DF-MP2 einsum example
+- [x] NKI fused contraction kernels (`mp2_energy`, `ao_to_mo_transform`)
+- [x] XLA operand residency (`to_xla` / `from_xla`)
+- [x] NKI CPU simulator + `nki-simulator` CI gate
+- [x] Optimal contraction ordering — greedy path search for 3+ operands
+- [x] Multi-contraction shared-operand XLA pre-pinning
+- [x] Contraction plan cache (`clear_plan_cache` / `plan_cache_info`)
+- [x] Alpha/beta scaling for `einsum` (cuTENSOR GEMM-style)
+- [x] Input validation with descriptive errors
+- [x] Tensor Train (TT) decomposition (TT-SVD)
+- [x] Non-negative CP + warm-start CP
+- [x] PEP 561 `py.typed` marker
+- [ ] Mixed precision / `dtype` override for bf16/fp16
+
+## Related Projects
+
+| Project | What |
+|---------|------|
+| [trnfft](https://github.com/scttfrdmn/trnfft) | FFT + complex ops |
+| [trnblas](https://github.com/scttfrdmn/trnblas) | BLAS operations |
+| [trnsolver](https://github.com/scttfrdmn/trnsolver) | Linear solvers |
+| [trnrand](https://github.com/scttfrdmn/trnrand) | Random number generation |
+| [trnsparse](https://github.com/scttfrdmn/trnsparse) | Sparse operations |
+
+## License
+
+Apache 2.0 — Copyright 2026 Scott Friedman
+
+
+## Disclaimer
+
+trnsci is an **independent open-source project**. It is not sponsored by, endorsed by, or affiliated with Amazon.com, Inc., Amazon Web Services, Inc., or Annapurna Labs Ltd.
+
+"AWS", "Amazon", "Trainium", "Inferentia", "NeuronCore", "Neuron SDK", and related identifiers are trademarks of their respective owners and are used here solely for descriptive and interoperability purposes. Use does not imply endorsement, partnership, or any other relationship.
+
+All work, opinions, analyses, benchmark results, architectural commentary, and editorial judgments in this repository and on [trnsci.dev](https://trnsci.dev) are those of the project's contributors. They do not represent the views, positions, or commitments of Amazon, AWS, or Annapurna Labs.
+
+Feedback directed at the Neuron SDK or Trainium hardware is good-faith ecosystem commentary from independent users. It is not privileged information, is not pre-reviewed by AWS, and should not be read as authoritative about product roadmap, behavior, or quality.
+
+For official AWS guidance, see [aws-neuron documentation](https://awsdocs-neuron.readthedocs-hosted.com/) and the [AWS Trainium product page](https://aws.amazon.com/ai/machine-learning/trainium/).

trntensor-0.3.0/README.md

@@ -0,0 +1,112 @@
+# trntensor
+
+[![CI](https://github.com/trnsci/trntensor/actions/workflows/ci.yml/badge.svg)](https://github.com/trnsci/trntensor/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/trnsci/trntensor/graph/badge.svg)](https://codecov.io/gh/trnsci/trntensor)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+[![PyPI](https://img.shields.io/pypi/v/trntensor)](https://pypi.org/project/trntensor/)
+[![Python](https://img.shields.io/pypi/pyversions/trntensor)](https://pypi.org/project/trntensor/)
+[![License](https://img.shields.io/github/license/trnsci/trntensor)](LICENSE)
+[![Docs](https://img.shields.io/badge/docs-trnsci.dev-blue)](https://trnsci.dev/trntensor/)
+
+Tensor contractions for AWS Trainium via NKI.
+
+Einstein summation with contraction planning, CP and Tucker decompositions. Expresses scientific tensor workloads naturally instead of decomposing to GEMM. Part of the trnsci scientific computing suite ([github.com/trnsci](https://github.com/trnsci)).
+
+## Current phase
+
+trntensor follows the [trnsci 5-phase roadmap](https://trnsci.dev/roadmap/). Active work is tracked in phase-labeled GitHub issues:
+
+- **[Phase 1 — correctness](https://github.com/trnsci/trntensor/issues/27)** (active): matmul + batched-matmul NKI kernels in place; awaiting hardware validation + additional `@pytest.mark.neuron` coverage.
+- **[Phase 2 — precision](https://github.com/trnsci/trntensor/issues/28)**: precision-aware contraction path selection (depends on [trnblas#22](https://github.com/trnsci/trnblas/issues/22) double-double GEMM).
+- **[Phase 3 — perf](https://github.com/trnsci/trntensor/issues/29)**: opt_einsum-style path planner, plan cache reuse.
+- **[Phase 4 — multi-chip](https://github.com/trnsci/trntensor/issues/30)**: sharded tensor contractions.
+- **[Phase 5 — generation](https://github.com/trnsci/trntensor/issues/31)**: trn2 fused multi-contraction paths.
+
+Suite-wide tracker: [trnsci/trnsci#1](https://github.com/trnsci/trnsci/issues/1).
+
+## Install
+
+```bash
+pip install trntensor
+```
+
+## Usage
+
+```python
+import torch
+import trntensor
+
+# Einsum — drop-in for torch.einsum with contraction planning
+C = trntensor.einsum("ij,jk->ik", A, B)           # matmul
+T = trntensor.einsum("ap,bp->ab", B_i, B_j)       # DF-MP2 contraction
+X = trntensor.einsum("mi,mnP->inP", C_occ, eri)   # AO→MO transform
+
+# Contraction planning
+plan = trntensor.plan_contraction("ij,jk->ik", A, B)
+flops = trntensor.estimate_flops("ij,jk->ik", A, B)
+
+# CP decomposition (tensor hypercontraction)
+factors, weights = trntensor.cp_decompose(tensor, rank=10)
+reconstructed = trntensor.cp_reconstruct(factors, weights)
+
+# Tucker decomposition (HOSVD)
+core, factors = trntensor.tucker_decompose(tensor, ranks=(5, 5, 5))
+```
+
+## Operations
+
+| Category | Operation | Description |
+|----------|-----------|-------------|
+| Contraction | `einsum` | General tensor contraction |
+| Contraction | `multi_einsum` | Multiple contractions (fusion-ready) |
+| Planning | `plan_contraction` | Analyze and select strategy |
+| Planning | `estimate_flops` | FLOPs for a contraction |
+| Decomposition | `cp_decompose` | CP/PARAFAC via ALS |
+| Decomposition | `tucker_decompose` | Tucker via HOSVD |
+
+## Status
+
+- [x] Einsum with matmul/bmm/torch dispatch
+- [x] Contraction planner
+- [x] CP decomposition (ALS)
+- [x] Tucker decomposition (HOSVD)
+- [x] DF-MP2 einsum example
+- [x] NKI fused contraction kernels (`mp2_energy`, `ao_to_mo_transform`)
+- [x] XLA operand residency (`to_xla` / `from_xla`)
+- [x] NKI CPU simulator + `nki-simulator` CI gate
+- [x] Optimal contraction ordering — greedy path search for 3+ operands
+- [x] Multi-contraction shared-operand XLA pre-pinning
+- [x] Contraction plan cache (`clear_plan_cache` / `plan_cache_info`)
+- [x] Alpha/beta scaling for `einsum` (cuTENSOR GEMM-style)
+- [x] Input validation with descriptive errors
+- [x] Tensor Train (TT) decomposition (TT-SVD)
+- [x] Non-negative CP + warm-start CP
+- [x] PEP 561 `py.typed` marker
+- [ ] Mixed precision / `dtype` override for bf16/fp16
+
+## Related Projects
+
+| Project | What |
+|---------|------|
+| [trnfft](https://github.com/scttfrdmn/trnfft) | FFT + complex ops |
+| [trnblas](https://github.com/scttfrdmn/trnblas) | BLAS operations |
+| [trnsolver](https://github.com/scttfrdmn/trnsolver) | Linear solvers |
+| [trnrand](https://github.com/scttfrdmn/trnrand) | Random number generation |
+| [trnsparse](https://github.com/scttfrdmn/trnsparse) | Sparse operations |
+
+## License
+
+Apache 2.0 — Copyright 2026 Scott Friedman
+
+
+## Disclaimer
+
+trnsci is an **independent open-source project**. It is not sponsored by, endorsed by, or affiliated with Amazon.com, Inc., Amazon Web Services, Inc., or Annapurna Labs Ltd.
+
+"AWS", "Amazon", "Trainium", "Inferentia", "NeuronCore", "Neuron SDK", and related identifiers are trademarks of their respective owners and are used here solely for descriptive and interoperability purposes. Use does not imply endorsement, partnership, or any other relationship.
+
+All work, opinions, analyses, benchmark results, architectural commentary, and editorial judgments in this repository and on [trnsci.dev](https://trnsci.dev) are those of the project's contributors. They do not represent the views, positions, or commitments of Amazon, AWS, or Annapurna Labs.
+
+Feedback directed at the Neuron SDK or Trainium hardware is good-faith ecosystem commentary from independent users. It is not privileged information, is not pre-reviewed by AWS, and should not be read as authoritative about product roadmap, behavior, or quality.
+
+For official AWS guidance, see [aws-neuron documentation](https://awsdocs-neuron.readthedocs-hosted.com/) and the [AWS Trainium product page](https://aws.amazon.com/ai/machine-learning/trainium/).