trntensor 0.1.1__tar.gz → 0.1.2__tar.gz

{trntensor-0.1.1 → trntensor-0.1.2}/.github/workflows/ci.yml

@@ -12,8 +12,8 @@ jobs:
       matrix:
         python-version: ["3.10", "3.11", "3.12"]
     steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-python@v5
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
       - run: pip install -e ".[dev]"

trntensor-0.1.2/.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}"

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

@@ -8,8 +8,8 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-python@v5
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
         with:
           python-version: "3.12"
       - run: pip install build
@@ -22,6 +22,7 @@ jobs:
   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
@@ -38,6 +39,7 @@ jobs:
   publish-testpypi:
     needs: build
     runs-on: ubuntu-latest
+    # Only publish pre-releases to TestPyPI
     if: ${{ contains(github.event.release.tag_name, '-') }}
     environment:
       name: testpypi

{trntensor-0.1.1 → trntensor-0.1.2}/CHANGELOG.md

@@ -5,6 +5,30 @@ 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).
 
+## [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

{trntensor-0.1.1 → trntensor-0.1.2}/CODE_OF_CONDUCT.md

@@ -12,7 +12,7 @@ This Code applies within all project spaces — the GitHub repositories under th
 
 Concerns about behaviour in project spaces may be reported confidentially to:
 
-**3011922+scttfrdmn@users.noreply.github.com**
+**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.
 

{trntensor-0.1.1/trntensor.egg-info → trntensor-0.1.2}/PKG-INFO

@@ -1,17 +1,18 @@
 Metadata-Version: 2.4
 Name: trntensor
-Version: 0.1.1
+Version: 0.1.2
 Summary: Tensor contractions for AWS Trainium via NKI
-Author-email: Scott Friedman <3011922+scttfrdmn@users.noreply.github.com>
+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.github.io/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

trntensor-0.1.2/benchmarks/bench_einsum.py

@@ -0,0 +1,56 @@
+"""Benchmarks for einsum dispatch and decompositions.
+
+Run with:
+    pytest benchmarks/ --benchmark-only
+
+On-hardware (via SSM orchestration):
+    scripts/run_benchmarks.sh trn1
+"""
+
+import pytest
+import torch
+
+import trntensor
+
+
+@pytest.fixture(autouse=True)
+def _seed():
+    torch.manual_seed(0)
+
+
+class TestEinsumBench:
+
+    def test_matmul_512(self, benchmark):
+        A = torch.randn(512, 512)
+        B = torch.randn(512, 512)
+        benchmark(trntensor.einsum, "ij,jk->ik", A, B)
+
+    def test_bmm_batched(self, benchmark):
+        A = torch.randn(16, 256, 256)
+        B = torch.randn(16, 256, 256)
+        benchmark(trntensor.einsum, "bij,bjk->bik", A, B)
+
+    def test_df_mp2_pair(self, benchmark):
+        """DF-MP2 pair contraction: T_ab = Σ_P B_ia^P B_jb^P."""
+        nvir, naux = 48, 128
+        Bi = torch.randn(nvir, naux)
+        Bj = torch.randn(nvir, naux)
+        benchmark(trntensor.einsum, "ap,bp->ab", Bi, Bj)
+
+    def test_4index_transform(self, benchmark):
+        """AO→MO half-transform."""
+        nbasis, nocc, naux = 32, 8, 64
+        C = torch.randn(nbasis, nocc)
+        eri = torch.randn(nbasis, nbasis, naux)
+        benchmark(trntensor.einsum, "mi,mnP->inP", C, eri)
+
+
+class TestDecomposeBench:
+
+    def test_cp_rank8(self, benchmark):
+        T = torch.randn(16, 16, 16)
+        benchmark(trntensor.cp_decompose, T, rank=8, max_iter=20)
+
+    def test_tucker_low_rank(self, benchmark):
+        T = torch.randn(16, 16, 16)
+        benchmark(trntensor.tucker_decompose, T, ranks=(4, 4, 4))

trntensor-0.1.2/docs/benchmarks.md

@@ -0,0 +1,38 @@
+# Benchmarks
+
+Performance results for trntensor — einsum dispatch, contraction planning, and CP / Tucker decompositions — comparing the PyTorch CPU fallback against the NKI Trainium path.
+
+## Reproducing
+
+```bash
+# CPU baseline (local)
+pytest benchmarks/bench_einsum.py --benchmark-only
+
+# On-hardware (via SSM orchestration)
+AWS_PROFILE=... ./scripts/run_benchmarks.sh trn1
+```
+
+## PyTorch CPU baseline
+
+Runs on a developer laptop. Use these numbers to sanity-check relative costs of different contraction patterns — not as a ceiling for what Trainium can do.
+
+_Platform: Darwin arm64, PyTorch CPU, median of ≥5 rounds._
+
+| Op | Shape | Median |
+|---|---|---|
+| `einsum ap,bp->ab` (DF-MP2 pair) | 48×128 × 48×128 | 4.5 µs |
+| `einsum mi,mnP->inP` (4-index) | 32×8, 32×32×64 | 10.5 µs |
+| `einsum ij,jk->ik` | 512×512 × 512×512 | 104 µs |
+| `tucker_decompose` | 16³ ranks (4,4,4) | 177 µs |
+| `einsum bij,bjk->bik` | 16×256×256 | 337 µs |
+| `cp_decompose` | 16³ rank 8 (20 iters) | 4.6 ms |
+
+## NKI (Trainium) — pending
+
+Fused-contraction kernels are scaffolded in `trntensor/nki/dispatch.py` but not yet implemented. Tracked in [#1](https://github.com/trnsci/trntensor/issues/1). The `scripts/run_benchmarks.sh` SSM orchestration is ready to populate the NKI column once kernels land.
+
+| Op | NKI (Trainium) | Speedup vs CPU |
+|---|---|---|
+| all | pending | pending |
+
+See [AWS Setup](aws_setup.md) for the on-hardware CI flow.

{trntensor-0.1.1 → trntensor-0.1.2}/infra/terraform/main.tf

@@ -119,7 +119,7 @@ resource "aws_instance" "ci" {
     #!/bin/bash
     set -euxo pipefail
     cd /home/ubuntu
-    sudo -u ubuntu git clone https://github.com/scttfrdmn/trntensor.git trntensor
+    sudo -u ubuntu git clone https://github.com/trnsci/trntensor.git trntensor
     # Install into the AMI's pre-built Neuron venv (has neuronxcc preinstalled).
     # Use [dev] only — [neuron] would try to fetch neuronxcc from PyPI where it doesn't exist.
     NEURON_VENV=$(ls -d /opt/aws_neuronx_venv_pytorch_* | head -1)

{trntensor-0.1.1 → trntensor-0.1.2}/mkdocs.yml

@@ -22,6 +22,7 @@ nav:
   - Quickstart: quickstart.md
   - API Reference: api.md
   - Architecture: architecture.md
+  - Benchmarks: benchmarks.md
   - Deployment:
       - AWS Setup for Neuron CI: aws_setup.md
 

{trntensor-0.1.1 → trntensor-0.1.2}/pyproject.toml

@@ -4,19 +4,20 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "trntensor"
-version = "0.1.1"
+version = "0.1.2"
 description = "Tensor contractions for AWS Trainium via NKI"
 readme = "README.md"
 license = "Apache-2.0"
 requires-python = ">=3.10"
 authors = [
-    { name = "Scott Friedman", email = "3011922+scttfrdmn@users.noreply.github.com" },
+    { name = "Scott Friedman", email = "scttfrdmn@gmail.com" },
 ]
 classifiers = [
     "Development Status :: 3 - Alpha",
     "Intended Audience :: Science/Research",
     "Programming Language :: Python :: 3",
     "Topic :: Scientific/Engineering",
+    "Topic :: Scientific/Engineering :: Mathematics",
 ]
 dependencies = [
     "torch>=2.1",
@@ -40,7 +41,7 @@ dev = [
 
 [project.urls]
 Homepage = "https://github.com/trnsci/trntensor"
-Documentation = "https://trnsci.github.io/trntensor/"
+Documentation = "https://trnsci.dev/trntensor/"
 Repository = "https://github.com/trnsci/trntensor"
 Issues = "https://github.com/trnsci/trntensor/issues"
 

{trntensor-0.1.1 → trntensor-0.1.2}/scripts/run_benchmarks.sh

@@ -81,7 +81,7 @@ BENCH_SCRIPT="source /opt/aws_neuronx_venv_pytorch_2_9/bin/activate && \
   git fetch --all && \
   git checkout $SHA && \
   pip install -e '.[dev]' --quiet && \
-  pytest benchmarks/bench_rand.py --benchmark-only --benchmark-json=$REMOTE_JSON --tb=short 2>&1 | tail -120 && \
+  pytest benchmarks/ --benchmark-only --benchmark-json=$REMOTE_JSON --tb=short 2>&1 | tail -120 && \
   echo BENCH_OK"
 
 CMD_ID=$(aws ssm send-command \

trntensor-0.1.2/tests/test_nki.py

@@ -0,0 +1,29 @@
+"""Backend dispatch tests (CPU path)."""
+
+import pytest
+import trntensor
+from trntensor.nki.dispatch import HAS_NKI
+
+
+def test_set_backend_auto():
+    trntensor.set_backend("auto")
+    assert trntensor.get_backend() == "auto"
+
+
+def test_set_backend_pytorch():
+    trntensor.set_backend("pytorch")
+    assert trntensor.get_backend() == "pytorch"
+    trntensor.set_backend("auto")
+
+
+def test_set_backend_invalid():
+    with pytest.raises(AssertionError):
+        trntensor.set_backend("tpu")
+
+
+@pytest.mark.skipif(HAS_NKI, reason="Test asserts the no-NKI guard")
+def test_set_backend_nki_raises_without_neuronxcc():
+    with pytest.raises(RuntimeError, match="neuronxcc"):
+        trntensor.set_backend("nki")
+    # State should not have changed
+    assert trntensor.get_backend() != "nki"

{trntensor-0.1.1 → trntensor-0.1.2}/trntensor/__init__.py

@@ -5,7 +5,7 @@ Einstein summation, contraction planning, and tensor decompositions
 for scientific computing. Part of the trnsci scientific computing suite.
 """
 
-__version__ = "0.1.1"
+__version__ = "0.1.2"
 
 from .einsum import einsum, multi_einsum
 from .plan import ContractionPlan, plan_contraction, estimate_flops

{trntensor-0.1.1 → trntensor-0.1.2}/trntensor/nki/dispatch.py

@@ -25,6 +25,10 @@ _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. Install with: pip install 'trntensor[neuron]'"
+        )
     _backend = backend
 
 def get_backend() -> str:

{trntensor-0.1.1 → trntensor-0.1.2/trntensor.egg-info}/PKG-INFO

@@ -1,17 +1,18 @@
 Metadata-Version: 2.4
 Name: trntensor
-Version: 0.1.1
+Version: 0.1.2
 Summary: Tensor contractions for AWS Trainium via NKI
-Author-email: Scott Friedman <3011922+scttfrdmn@users.noreply.github.com>
+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.github.io/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

{trntensor-0.1.1 → trntensor-0.1.2}/trntensor.egg-info/SOURCES.txt

@@ -8,11 +8,14 @@ README.md
 mkdocs.yml
 pyproject.toml
 .github/workflows/ci.yml
-.github/workflows/docs.yml
+.github/workflows/notify-umbrella.yml
 .github/workflows/publish.yml
+benchmarks/__init__.py
+benchmarks/bench_einsum.py
 docs/api.md
 docs/architecture.md
 docs/aws_setup.md
+docs/benchmarks.md
 docs/index.md
 docs/installation.md
 docs/quickstart.md
@@ -25,6 +28,7 @@ scripts/run_neuron_tests.sh
 tests/conftest.py
 tests/test_decompose.py
 tests/test_einsum.py
+tests/test_nki.py
 tests/test_plan.py
 trntensor/__init__.py
 trntensor/decompose.py

trntensor-0.1.1/.github/workflows/docs.yml

@@ -1,38 +0,0 @@
-name: Deploy Docs
-
-on:
-  push:
-    branches: [main]
-
-permissions:
-  contents: read
-  pages: write
-  id-token: write
-
-concurrency:
-  group: pages
-  cancel-in-progress: true
-
-jobs:
-  build:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-python@v5
-        with:
-          python-version: "3.12"
-      - run: pip install mkdocs mkdocs-material
-      - run: mkdocs build --strict
-      - uses: actions/upload-pages-artifact@v3
-        with:
-          path: site/
-
-  deploy:
-    needs: build
-    runs-on: ubuntu-latest
-    environment:
-      name: github-pages
-      url: ${{ steps.deployment.outputs.page_url }}
-    steps:
-      - id: deployment
-        uses: actions/deploy-pages@v4