Scientific computing suite for AWS Trainium via NKI. Six focused libraries providing the cu* equivalents the Neuron SDK doesn't ship.
| Project | NVIDIA analog | Scope |
|---|---|---|
| trnfft | cuFFT | FFT, complex tensors, complex NN layers |
| trnblas | cuBLAS | Level 1–3 BLAS, batched GEMM |
| trnrand | cuRAND | Philox PRNG, Sobol/Halton QMC, LHS |
| trnsolver | cuSOLVER | Cholesky/LU/QR, Jacobi eigh, CG/GMRES |
| trnsparse | cuSPARSE | CSR/COO, SpMV/SpMM, Schwarz screening |
| trntensor | cuTENSOR | einsum with contraction planning, CP/Tucker |
# Meta-package with all six sub-projects
pip install trnsci[all]
# Individual components
pip install trnsci[fft] # just trnfft
pip install trnsci[blas] # just trnblas
# ... etc
# On Neuron hardware
pip install trnsci[all,neuron]git clone [email protected]:trnsci/trnsci.git
cd trnsci
make install-dev # pip install -e on each sub-project + umbrella
make test-all # run pytest across all sub-projectspython examples/quantum_chemistry/df_mp2_synthetic.py --demoDF-MP2 energy evaluation composing trnblas (half-transforms), trnsolver (Cholesky of DF metric), and trntensor (einsum contraction). See also examples/nvidia_samples/ for direct ports of canonical NVIDIA CUDA samples.
All sub-projects are Alpha. CPU/PyTorch fallback is functional end-to-end. NKI kernels are scaffolded across the suite; on-hardware validation is the next milestone.
Five-phase plan from current alpha to generation-tuned stable release. See ROADMAP.md or the reader-oriented version on trnsci.dev.
Apache 2.0 — Copyright 2026 Scott Friedman.
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.