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).

Current phase

trntensor follows the trnsci 5-phase roadmap. Active work is tracked in phase-labeled GitHub issues:

• Phase 1 — correctness (active): matmul + batched-matmul NKI kernels in place; awaiting hardware validation + additional @pytest.mark.neuron coverage.
• Phase 2 — precision: precision-aware contraction path selection (depends on trnblas#22 double-double GEMM).
• Phase 3 — perf: opt_einsum-style path planner, plan cache reuse.
• Phase 4 — multi-chip: sharded tensor contractions.
• Phase 5 — generation: trn2 fused multi-contraction paths.

Suite-wide tracker: trnsci/trnsci#1.

Install

pip install trntensor

# Trn2/Trn3 with TorchNeuron (native PyTorch, Neuron SDK 2.33+)
pip install "trntensor[torchneuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com

# Trn1 / PyTorch-XLA lineage
pip install "trntensor[neuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com

Kernels dispatch on whichever Neuron PyTorch lineage is installed:

• TorchNeuron (native): Trn2/Trn3, Neuron SDK 2.33+, Python ≥ 3.12, torch 2.14, torch-neuronx 2.14.*. Each kernel is a PyTorch custom op (torch.ops.trntensor.*) on the neuron device; no torch_xla.
• PyTorch/XLA: Trn1 (and any host on the torch-neuronx 2.9 lineage). Kernels dispatch through torch_xla, unchanged.

trntensor picks the native path when the native torch_neuronx and a neuron device are present, and the XLA path when torch_xla is installed; otherwise it runs on PyTorch. Set TRNSCI_NKI_DISPATCH=native or xla (shared by all trnsci libraries) to force a path. Pin operands on the device with to_device / from_device; the older to_xla / from_xla still work on both lineages (see the migration note). With an editable install (pip install -e), set NKI_ENABLE_TRACE_CACHE=0: the NKI kernel cache only versions installed wheels and can serve stale kernels.

Usage

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

Status

• Einsum with matmul/bmm/torch dispatch
• Contraction planner
• CP decomposition (ALS)
• Tucker decomposition (HOSVD)
• DF-MP2 einsum example
• NKI fused contraction kernels (mp2_energy, ao_to_mo_transform)
• Device operand residency (to_device / from_device; to_xla / from_xla kept)
• Native TorchNeuron dispatch (Trn2/Trn3, torch.ops.trntensor.*) alongside PyTorch/XLA (Trn1)
• NKI CPU simulator + nki-simulator CI gate
• Optimal contraction ordering — greedy path search for 3+ operands
• Multi-contraction shared-operand XLA pre-pinning
• Contraction plan cache (clear_plan_cache / plan_cache_info)
• Alpha/beta scaling for einsum (cuTENSOR GEMM-style)
• Input validation with descriptive errors
• Tensor Train (TT) decomposition (TT-SVD)
• Non-negative CP + warm-start CP
• PEP 561 py.typed marker
• Mixed precision / dtype override for bf16/fp16
• precision= kwarg ("fast" / "kahan" fp64 promotion / "dd" stub)

Related Projects

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 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 and the AWS Trainium product page.