trnsparse 0.3.0__tar.gz → 0.3.2__tar.gz

trnsparse-0.3.2/.github/workflows/ci.yml

@@ -0,0 +1,58 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - run: pip install "ruff>=0.6"
+      - run: ruff check .
+      - run: ruff format --check .
+
+  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 and not nki_simulator" --cov=trnsparse --cov-report=xml
+      - name: Upload coverage reports to Codecov
+        if: matrix.python-version == '3.12'
+        uses: codecov/codecov-action@v5
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          slug: trnsci/trnsparse
+
+  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: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - name: Install trnsparse + NKI simulator deps
+        run: |
+          pip install -e ".[dev]"
+          pip install --extra-index-url https://pip.repos.neuron.amazonaws.com \
+            "nki>=0.3.0"
+      - name: Run simulator-backed kernel tests
+        env:
+          TRNSPARSE_USE_SIMULATOR: "1"
+        run: pytest tests/ -v -m nki_simulator --tb=short

trnsparse-0.3.2/.pre-commit-config.yaml

@@ -0,0 +1,7 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.8.6
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format

{trnsparse-0.3.0 → trnsparse-0.3.2}/CHANGELOG.md

@@ -5,6 +5,73 @@ 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.3.2] — 2026-04-14
+
+### Added
+
+- **`cg_bsr`** and **`power_iteration_bsr`** — Conjugate Gradient and
+  power iteration on block-sparse row matrices. Plumbing on top of
+  `bsr_spmm` (one kernel dispatch per iteration). Closes Phase 1 of
+  #22 on-chip iterative solvers.
+- **`jacobi_preconditioner_bsr(A)`** — builds a diagonal preconditioner
+  for `cg_bsr`'s `M=` argument.
+- **`bsr_diagonal(A)`** — extracts the main diagonal from a BSR matrix.
+- **`docs/iterative_solvers.md`** — design note covering the v0.3.2
+  plumbing and the v0.4.0 fused-kernel goal (#24). Explains the
+  architectural win Trainium offers (A SBUF-resident across iterations)
+  vs the current per-iteration HBM round-trip.
+- **`tests/test_iterative.py`** — 8 CPU tests including scipy parity at
+  `atol=1e-4` on a 128×128 SPD system.
+- **`benchmarks/bench_iterative.py`** — cg_bsr vs scipy.sparse.linalg.cg.
+  At 128×128 SPD: scipy 310 μs, trnsparse 369 μs (1.19×).
+
+### Notes
+
+- Algorithm body for CG is a local copy of `trnsolver.iterative.cg`;
+  kept local to avoid a cross-repo runtime dependency for one function.
+- v0.4.0 will layer the fused CG/power-iteration NKI kernel on top —
+  tracked in #24. The API stays stable across the transition; users
+  upgrading from v0.3.2 get the fused-kernel speedup automatically
+  when the fused path is available.
+
+## [0.3.1] — 2026-04-14
+
+### Changed
+
+- **Migrated NKI imports to the `nki.*` namespace** (NKI 0.3.0 Stable,
+  Neuron SDK 2.29, April 2026). Legacy `neuronxcc.nki.*` shim is no
+  longer used. `pyproject.toml` `[neuron]` extra gains `nki>=0.3.0`
+  alongside the existing `neuronxcc>=2.24` and `torch-neuronx>=2.9`.
+  Hosts without an `nki` wheel (macOS, non-Linux archs) still hit
+  `HAS_NKI=False` and get the torch fallback. Kernel bodies unchanged —
+  the trnblas audit confirmed the positional `nisa.nc_matmul` +
+  `nl.copy(psum, ...)` pattern complies with NKI 0.3.0.
+- `test` CI job now filters `-m "not neuron and not nki_simulator"` so
+  each test runs in exactly one job.
+
+### Added
+
+- **`TRNSPARSE_USE_SIMULATOR=1` dispatch branch** through
+  `nki.simulate(kernel)(np_args)`. Bypasses torch_xla + NEFF compile;
+  kernels run on CPU for correctness iteration. Hardware still owns
+  perf numbers.
+- **`nki-simulator` CI job on `ubuntu-latest`** — installs `nki>=0.3.0`
+  from the AWS pip index and runs the simulator suite on every push/PR.
+  Kernel correctness gate without AWS cost. Catches Python-trace-level
+  errors (bad kwargs, dropped ops, shape mismatches); MLIR verifier
+  errors remain hardware-only (NKI 0.3.0 has no documented device-free
+  NEFF compile API).
+- `tests/test_nki_sim.py` — curated simulator suite (4 tests: CSR
+  aligned + rectangular, BSR block-dense + block-diagonal). Skips
+  cleanly off-hardware.
+- `scripts/run_simulator_tests.sh` — SSM runner mirroring
+  `run_neuron_tests.sh` with `TRNSPARSE_USE_SIMULATOR=1` in the env.
+- `tests/conftest.py` — registers the `nki_simulator` pytest marker.
+
+Addresses [trnsci/trnsparse#23](https://github.com/trnsci/trnsparse/issues/23).
+Follows the trnblas reference commits `c693561`, `f24993b`, `77eeb82`
+(suite-wide coordination in `trnsci/trnsci#5`).
+
 ## [0.3.0] — 2026-04-13
 
 ### Added

{trnsparse-0.3.0 → trnsparse-0.3.2}/CLAUDE.md

@@ -64,3 +64,15 @@ pip install -e ".[dev]"
 pytest tests/ -v
 python examples/sparse_fock.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.

{trnsparse-0.3.0/trnsparse.egg-info → trnsparse-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trnsparse
-Version: 0.3.0
+Version: 0.3.2
 Summary: Sparse matrix operations for AWS Trainium via NKI
 Author-email: Scott Friedman <scttfrdmn@gmail.com>
 License-Expression: Apache-2.0
@@ -19,17 +19,21 @@ License-File: LICENSE
 Requires-Dist: torch>=2.1
 Requires-Dist: numpy>=1.24
 Provides-Extra: neuron
+Requires-Dist: nki>=0.3.0; 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: pytest-cov>=4.1; extra == "dev"
 Requires-Dist: scipy>=1.11; extra == "dev"
 Dynamic: license-file
 
 # trnsparse
 
 [![CI](https://github.com/trnsci/trnsparse/actions/workflows/ci.yml/badge.svg)](https://github.com/trnsci/trnsparse/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/trnsci/trnsparse/graph/badge.svg)](https://codecov.io/gh/trnsci/trnsparse)
+[![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/trnsparse)](https://pypi.org/project/trnsparse/)
 [![Python](https://img.shields.io/pypi/pyversions/trnsparse)](https://pypi.org/project/trnsparse/)
 [![License](https://img.shields.io/github/license/trnsci/trnsparse)](LICENSE)

{trnsparse-0.3.0 → trnsparse-0.3.2}/README.md

@@ -1,6 +1,8 @@
 # trnsparse
 
 [![CI](https://github.com/trnsci/trnsparse/actions/workflows/ci.yml/badge.svg)](https://github.com/trnsci/trnsparse/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/trnsci/trnsparse/graph/badge.svg)](https://codecov.io/gh/trnsci/trnsparse)
+[![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/trnsparse)](https://pypi.org/project/trnsparse/)
 [![Python](https://img.shields.io/pypi/pyversions/trnsparse)](https://pypi.org/project/trnsparse/)
 [![License](https://img.shields.io/github/license/trnsci/trnsparse)](LICENSE)

{trnsparse-0.3.0 → trnsparse-0.3.2}/benchmarks/bench_bsr_spmm.py

@@ -21,8 +21,7 @@ import torch
 import trnsparse
 from trnsparse.nki.dispatch import HAS_NKI
 
-
-M_BLOCKS = [4, 8]         # matrix has M_BLOCKS * 128 rows
+M_BLOCKS = [4, 8]  # matrix has M_BLOCKS * 128 rows
 N_BLOCKS = [4, 8]
 BLOCK_DENSITIES = [0.1, 0.25, 0.5]
 RHS_COLS = [128, 256]
@@ -58,7 +57,7 @@ def bsr_and_B(m_blocks, n_blocks, block_density, bsr_rhs_cols):
     for i in range(m_blocks):
         for j in range(n_blocks):
             if mask[i, j]:
-                A[i * b:(i + 1) * b, j * b:(j + 1) * b] = torch.randn(b, b)
+                A[i * b : (i + 1) * b, j * b : (j + 1) * b] = torch.randn(b, b)
     bsr = trnsparse.BSRMatrix.from_dense(A, block_size=b)
     B = torch.randn(N, bsr_rhs_cols)
     return bsr, B, A

trnsparse-0.3.2/benchmarks/bench_iterative.py

@@ -0,0 +1,50 @@
+"""Iterative-solver benchmarks: trnsparse.cg_bsr vs scipy baseline.
+
+The v0.3.2 plumbing dispatches one `bsr_spmm` call per CG iteration.
+On NKI that means one kernel launch + HBM round-trip per iteration.
+Expect the current path to be dominated by dispatch overhead compared
+to scipy's compiled C loop, which motivates the v0.4.0 fused-kernel
+follow-up.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+import trnsparse
+
+
+@pytest.fixture(params=[128, 256])
+def iter_size(request):
+    return request.param
+
+
+@pytest.fixture
+def spd_bsr_and_dense(iter_size):
+    torch.manual_seed(0)
+    n = iter_size
+    M = torch.randn(n, n)
+    A_dense = M @ M.T + n * torch.eye(n)
+    A_bsr = trnsparse.BSRMatrix.from_dense(A_dense, block_size=128)
+    b = torch.randn(n)
+    return A_dense, A_bsr, b
+
+
+def test_cg_bsr_trnsparse(benchmark, spd_bsr_and_dense):
+    _, A_bsr, b = spd_bsr_and_dense
+    benchmark(lambda: trnsparse.cg_bsr(A_bsr, b, tol=1e-8, max_iter=2 * b.shape[0]))
+
+
+def test_cg_scipy(benchmark, spd_bsr_and_dense):
+    sp = pytest.importorskip("scipy.sparse")
+    spla = pytest.importorskip("scipy.sparse.linalg")
+    A_dense, _, b = spd_bsr_and_dense
+    A_scipy = sp.csr_matrix(A_dense.numpy())
+    b_np = b.numpy()
+    benchmark(lambda: spla.cg(A_scipy, b_np, rtol=1e-8, maxiter=2 * len(b_np)))
+
+
+def test_power_iteration_trnsparse(benchmark, spd_bsr_and_dense):
+    _, A_bsr, _ = spd_bsr_and_dense
+    benchmark(lambda: trnsparse.power_iteration_bsr(A_bsr, max_iter=500, tol=1e-9))

{trnsparse-0.3.0 → trnsparse-0.3.2}/benchmarks/bench_spmm.py

@@ -14,6 +14,7 @@ On Trainium + neuronxcc installed, all four run in one pytest invocation.
 """
 
 import pytest
+
 import trnsparse
 from trnsparse.nki.dispatch import HAS_NKI
 

trnsparse-0.3.2/docs/iterative_solvers.md

@@ -0,0 +1,115 @@
+# Iterative solvers over BSR
+
+trnsparse v0.3.2 adds `cg_bsr` and `power_iteration_bsr` — Conjugate
+Gradient and power iteration on block-sparse row matrices. The API is
+stable; the architectural story below explains why there's a v0.4.0
+follow-up.
+
+## Why this matters
+
+Large SPD linear systems and dominant-eigenpair problems show up
+across scientific computing:
+
+- **Quantum chemistry**: Hamiltonian eigenvalue problems (HF, DFT,
+  CI), response equations (CPSCF).
+- **PDE discretizations**: stiffness-matrix solves for finite element
+  methods, graph Laplacian systems.
+- **Graph learning**: spectral embeddings, PageRank-like iterations.
+
+The matrix `A` in each case is typically block-sparse (Fock matrices
+after Schwarz screening; FEM stiffness tied to mesh connectivity;
+graph adjacency). BSR is the Trainium-native representation for those
+matrices (see `architecture.md`).
+
+## v0.3.2 — plumbing
+
+```python
+import trnsparse
+
+A = trnsparse.BSRMatrix.from_dense(fock_matrix, block_size=128)
+b = compute_rhs()
+
+x, iters, rel = trnsparse.cg_bsr(A, b, tol=1e-6, max_iter=1000)
+# Jacobi-preconditioned variant:
+M = trnsparse.jacobi_preconditioner_bsr(A)
+x, iters, rel = trnsparse.cg_bsr(A, b, tol=1e-6, M=M)
+
+lam, v, iters = trnsparse.power_iteration_bsr(A, max_iter=500)
+```
+
+Under the hood, each CG iteration calls `bsr_spmm(A, x.unsqueeze(1))`
+once. On the NKI backend that's one kernel dispatch + one HBM
+round-trip per iteration. On CPU it's `torch.sparse`-backed and
+roughly on par with `scipy.sparse.linalg.cg` (benchmarked: 369 μs vs
+310 μs at 128×128 SPD, 1.19× slower).
+
+## v0.4.0 — fused kernel with SBUF-resident A
+
+The architectural claim from
+[#22](https://github.com/trnsci/trnsparse/issues/22): Trainium's 32 GB
+SBUF per NeuronCore fits a 5000×5000 BSR Hamiltonian on-chip. CG
+doesn't need to round-trip `A` to HBM at all — only `x`, `r`, and `p`.
+
+The shape of the v0.4.0 kernel:
+
+```python
+@nki.jit
+def _cg_spd_kernel(A_blocks, A_cols, A_row_ptrs, b, max_iter):
+    # Load A blocks once into SBUF at the top.
+    A_sbuf = nl.load(A_blocks)
+
+    # State: x, r, p in SBUF registers.
+    x = nl.zeros(...)
+    r = nl.copy(b)
+    p = nl.copy(r)
+    rr = nl.reduce(r * r)
+
+    for k in nl.affine_range(max_iter):
+        Ap = _bsr_matvec_sbuf(A_sbuf, A_cols, A_row_ptrs, p)   # all SBUF
+        pAp = nl.reduce(p * Ap)
+        alpha = rr / pAp
+        x = x + alpha * p
+        r = r - alpha * Ap
+        rr_new = nl.reduce(r * r)
+        beta = rr_new / rr
+        p = r + beta * p
+        rr = rr_new
+
+    return x, residual_norm_history
+```
+
+Fixed `max_iter`, no early exit — returning the full residual history
+lets the host pick the convergence point post-hoc (standard
+NKI-inside-loop constraint: no dynamic control flow).
+
+**Expected regime where this wins**: when
+`max_iter × dispatch_overhead > fused_kernel_cost + hbm_load_A_once`.
+For a 4000×4000 BSR Hamiltonian with ~100 iterations to convergence,
+that's roughly an order of magnitude of wall-time reduction.
+
+**What needs to land first**:
+
+1. Simulator-iterated kernel skeleton (now tractable thanks to the
+   `nki-simulator` CI gate — see the NKI 0.3.0 migration in v0.3.1).
+2. SBUF sizing model for `A` — some workloads overflow 32 GB; then
+   fall back to the v0.3.2 plumbing.
+3. Dispatcher logic in `cg_bsr` that picks the fused kernel when
+   `max_iter * n` exceeds a threshold.
+
+Tracked in a dedicated sub-issue on `trnsci/trnsparse`.
+
+## Why CG and power iteration, not GMRES / Lanczos / Davidson
+
+v0.3.2 covers the two most common algorithm families:
+
+- **CG** for SPD systems — the workhorse for Hamiltonian solves,
+  stiffness systems, and many PDE discretizations.
+- **Power iteration** for dominant eigenpairs — the starting point
+  for spectral methods, PageRank-like fixed points, and the iteration
+  core inside Lanczos / Arnoldi / Davidson.
+
+GMRES (general non-symmetric systems), Lanczos / Arnoldi (full
+spectrum), and Davidson (interior eigenvalues) are follow-ups when
+users ask for them. The v0.4.0 fused kernel's structure (load A once,
+iterate on-chip) generalizes trivially to those variants — it's the
+same architectural pattern.

{trnsparse-0.3.0 → trnsparse-0.3.2}/examples/sparse_fock.py

@@ -12,7 +12,9 @@ Usage:
 
 import argparse
 import time
+
 import torch
+
 import trnsparse
 
 
@@ -27,7 +29,7 @@ def main():
         args.nbasis = 50
 
     n = args.nbasis
-    print(f"Sparse Fock build:")
+    print("Sparse Fock build:")
     print(f"  Basis functions: {n}")
     print(f"  Threshold:       {args.threshold:.0e}")
 
@@ -40,7 +42,7 @@ def main():
 
     # Screen
     stats = trnsparse.sparsity_stats(Q, args.threshold)
-    print(f"\n  Sparsity statistics:")
+    print("\n  Sparsity statistics:")
     print(f"    Total shell pairs:       {stats['total_pairs']}")
     print(f"    Significant pairs:       {stats['significant_pairs']}")
     print(f"    Pair sparsity:           {stats['pair_sparsity']:.1%}")
@@ -51,7 +53,7 @@ def main():
     integrals_dense = torch.randn(n, n) * Q * 0.01
     integrals_dense[~mask] = 0.0
     integrals_sparse = trnsparse.from_dense(integrals_dense)
-    print(f"    Integral matrix nnz:     {integrals_sparse.nnz} / {n*n}")
+    print(f"    Integral matrix nnz:     {integrals_sparse.nnz} / {n * n}")
 
     # Density matrix (random SPD for demo)
     P = torch.randn(n, n) * 0.1

trnsparse-0.3.2/pyproject.toml

@@ -0,0 +1,84 @@
+[build-system]
+requires = ["setuptools>=68.0", "setuptools-scm>=8.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "trnsparse"
+version = "0.3.2"
+description = "Sparse matrix operations for AWS Trainium via NKI"
+readme = "README.md"
+license = "Apache-2.0"
+requires-python = ">=3.10"
+authors = [
+    { 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",
+    "numpy>=1.24",
+]
+
+[project.optional-dependencies]
+neuron = [
+    # NKI 0.3.0 Stable (Neuron SDK 2.29+, April 2026). trnsparse imports
+    # from the canonical `nki.*` namespace; the legacy `neuronxcc.nki.*`
+    # shim is not used. SDK 2.29 DLAMI bundles nki==0.3.0 alongside
+    # neuronxcc==2.24.5133. Hosts without an `nki` wheel (macOS,
+    # non-Linux archs) still hit HAS_NKI=False and get the torch fallback.
+    "nki>=0.3.0",
+    "neuronxcc>=2.24",
+    "torch-neuronx>=2.9",
+]
+dev = ["pytest>=7.0", "pytest-benchmark>=4.0", "pytest-cov>=4.1", "scipy>=1.11"]
+
+[project.urls]
+Homepage = "https://github.com/trnsci/trnsparse"
+Documentation = "https://trnsci.dev/trnsparse/"
+Repository = "https://github.com/trnsci/trnsparse"
+Issues = "https://github.com/trnsci/trnsparse/issues"
+
+[tool.setuptools.packages.find]
+include = ["trnsparse*"]
+
+[tool.ruff]
+# Match the rest of the trnsci suite. Line length 100 is a compromise
+# between black's 88 and the wider tables common in scientific code.
+line-length = 100
+target-version = "py310"
+extend-exclude = ["site", "infra/terraform*/.terraform"]
+
+[tool.ruff.lint]
+# Sensible default selection: pycodestyle (E/W), Pyflakes (F), isort (I),
+# pyupgrade (UP), flake8-bugbear (B), flake8-simplify (SIM). Skip docstring
+# rules (D*) — we're deliberately light on docstrings in this project.
+select = ["E", "W", "F", "I", "UP", "B", "SIM"]
+ignore = [
+    "E501",   # line too long — formatter handles it
+    "B008",   # function call in default arg — common + intentional here
+    "SIM108", # if/else over ternary — readability call, not worth fighting
+    "SIM300", # Yoda-condition false positives on array comparisons
+    "E741",   # ambiguous single-letter names — `I` for identity matrix is idiomatic in linear algebra
+]
+
+[tool.ruff.lint.per-file-ignores]
+# NKI kernels use patterns that trip some lint rules (SBUF writes that
+# look unused to the linter, etc). Don't fight the kernel authoring loop.
+"trnsparse/nki/dispatch.py" = ["F841", "B007"]
+"trnsparse/nki/kernels.py"  = ["F841", "B007"]
+"scripts/*" = ["F401"]
+"tests/*"   = ["F401", "F811"]
+
+[tool.ruff.format]
+# Black-compatible defaults.
+quote-style = "double"
+indent-style = "space"
+
+[tool.pytest.ini_options]
+markers = ["neuron: requires Neuron hardware"]
+testpaths = ["tests"]

{trnsparse-0.3.0 → trnsparse-0.3.2}/scripts/bench_to_md.py

@@ -67,7 +67,7 @@ def render_markdown(rows: dict, machine_info: dict | None = None) -> str:
     out.append("| Operation | Variant | Param | Median (μs) | vs trnsparse-PyTorch |")
     out.append("|-----------|---------|-------|------------:|-------------------:|")
 
-    for (group, op, param) in sorted(rows.keys()):
+    for group, op, param in sorted(rows.keys()):
         variants = rows[(group, op, param)]
         baseline = variants.get("trnrand_pytorch")
         for variant in ("nki", "trnrand_pytorch", "torch"):
@@ -82,11 +82,9 @@ def render_markdown(rows: dict, machine_info: dict | None = None) -> str:
                 if ratio >= 1.0:
                     speedup = f"{ratio:.2f}× faster"
                 else:
-                    speedup = f"{1/ratio:.2f}× slower"
+                    speedup = f"{1 / ratio:.2f}× slower"
             param_disp = param.replace("_", " ") if param else "-"
-            out.append(
-                f"| {group}.{op} | {label} | {param_disp} | {us:>10.2f} | {speedup} |"
-            )
+            out.append(f"| {group}.{op} | {label} | {param_disp} | {us:>10.2f} | {speedup} |")
     return "\n".join(out) + "\n"
 
 

trnsparse-0.3.2/scripts/run_simulator_tests.sh

@@ -0,0 +1,116 @@
+#!/usr/bin/env bash
+#
+# Run NKI simulator-backed tests on the trn1 CI instance.
+#
+# Simulator bypasses torch_xla + NEFF compile: kernels run on CPU via
+# nki.simulate(kernel)(numpy_args). Use for correctness / constraint
+# iteration; hardware still owns perf numbers.
+#
+# Usage:
+#   AWS_PROFILE=aws ./scripts/run_simulator_tests.sh
+#
+# Same trap-stop pattern as run_neuron_tests.sh. Runs the nki_simulator-
+# marked test suite with TRNSPARSE_USE_SIMULATOR=1 set in the SSM env.
+#
+# Still AWS-resident for now (the nki wheel is linux_x86_64 only + lives
+# on the AWS pip index, not a common macOS target). The CI job on
+# ubuntu-latest covers the same surface on every push.
+
+set -euo pipefail
+
+INSTANCE_TYPE="${INSTANCE_TYPE:-trn1}"
+TAG="trnsparse-ci-${INSTANCE_TYPE}"
+REGION="${AWS_REGION:-us-east-1}"
+SHA="$(git rev-parse HEAD)"
+
+: "${AWS_PROFILE:?Set AWS_PROFILE, e.g. AWS_PROFILE=aws ./scripts/run_simulator_tests.sh}"
+
+echo "Looking up instance with Name=$TAG in $REGION..."
+INSTANCE_ID=$(aws ec2 describe-instances \
+  --filters "Name=tag:Name,Values=$TAG" \
+            "Name=instance-state-name,Values=stopped,running,pending" \
+  --query 'Reservations[0].Instances[0].InstanceId' \
+  --output text \
+  --region "$REGION")
+
+if [[ -z "$INSTANCE_ID" || "$INSTANCE_ID" == "None" ]]; then
+  echo "ERROR: No instance found with Name=$TAG" >&2
+  exit 1
+fi
+echo "Instance: $INSTANCE_ID"
+
+cleanup() {
+  local exit_code=$?
+  echo ""
+  echo "Stopping $INSTANCE_ID..."
+  aws ec2 stop-instances --instance-ids "$INSTANCE_ID" --region "$REGION" >/dev/null
+  exit "$exit_code"
+}
+trap cleanup EXIT
+
+STATE=$(aws ec2 describe-instances --instance-ids "$INSTANCE_ID" --region "$REGION" \
+  --query 'Reservations[0].Instances[0].State.Name' --output text)
+
+if [[ "$STATE" == "stopped" ]]; then
+  echo "Starting instance..."
+  aws ec2 start-instances --instance-ids "$INSTANCE_ID" --region "$REGION" >/dev/null
+fi
+
+echo "Waiting for instance-running..."
+aws ec2 wait instance-running --instance-ids "$INSTANCE_ID" --region "$REGION"
+echo "Waiting for SSM agent..."
+for _ in $(seq 1 60); do
+  PING=$(aws ssm describe-instance-information \
+    --filters "Key=InstanceIds,Values=$INSTANCE_ID" \
+    --region "$REGION" \
+    --query 'InstanceInformationList[0].PingStatus' --output text 2>/dev/null || true)
+  [[ "$PING" == "Online" ]] && break
+  sleep 5
+done
+if [[ "$PING" != "Online" ]]; then
+  echo "ERROR: SSM agent not Online after 5 minutes (last PingStatus=$PING)" >&2
+  exit 1
+fi
+
+echo "Sending simulator test command (SHA=$SHA)..."
+CMD_ID=$(aws ssm send-command \
+  --instance-ids "$INSTANCE_ID" \
+  --document-name "AWS-RunShellScript" \
+  --comment "trnsparse nki simulator tests @ $SHA" \
+  --parameters "commands=[
+    \"bash -c 'set -euo pipefail; cd /home/ubuntu/trnsparse && sudo -u ubuntu git fetch --all && sudo -u ubuntu git checkout $SHA && NEURON_VENV=\$(ls -d /opt/aws_neuronx_venv_pytorch_* | head -1) && sudo -u ubuntu \$NEURON_VENV/bin/pip install -e /home/ubuntu/trnsparse[dev] --quiet && sudo -u ubuntu env PATH=\$NEURON_VENV/bin:/usr/bin:/bin TRNSPARSE_USE_SIMULATOR=1 \$NEURON_VENV/bin/pytest /home/ubuntu/trnsparse/tests/ -v -m nki_simulator --tb=short'\"
+  ]" \
+  --region "$REGION" \
+  --output text --query 'Command.CommandId')
+
+echo "Command ID: $CMD_ID"
+echo "Waiting for command to complete..."
+for _ in $(seq 1 60); do
+  STATUS=$(aws ssm get-command-invocation \
+    --command-id "$CMD_ID" \
+    --instance-id "$INSTANCE_ID" \
+    --region "$REGION" \
+    --query 'Status' --output text 2>/dev/null || echo "InProgress")
+  [[ "$STATUS" != "InProgress" && "$STATUS" != "Pending" ]] && break
+  sleep 15
+done
+
+echo ""
+echo "=== STDOUT ==="
+aws ssm get-command-invocation \
+  --command-id "$CMD_ID" \
+  --instance-id "$INSTANCE_ID" \
+  --region "$REGION" \
+  --query 'StandardOutputContent' --output text
+
+echo ""
+echo "=== STDERR ==="
+aws ssm get-command-invocation \
+  --command-id "$CMD_ID" \
+  --instance-id "$INSTANCE_ID" \
+  --region "$REGION" \
+  --query 'StandardErrorContent' --output text | head -20
+
+echo ""
+echo "=== Status: $STATUS ==="
+[[ "$STATUS" == "Success" ]]

trnsparse-0.3.2/tests/conftest.py

@@ -0,0 +1,12 @@
+"""Test configuration."""
+
+import pytest  # noqa: F401 — imported for marker side-effect exposure
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "neuron: requires Neuron hardware")
+    config.addinivalue_line(
+        "markers",
+        "nki_simulator: runs NKI kernels via nki.simulate on CPU "
+        "(requires TRNSPARSE_USE_SIMULATOR=1 + nki>=0.3.0)",
+    )