Skip to content

trnsci/trnsci

Repository files navigation

trnsci

CI codecov pre-commit PyPI Python License Docs

Scientific computing suite for AWS Trainium via NKI. Six focused libraries providing the cu* equivalents the Neuron SDK doesn't ship.

Suite

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

Install

# 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]

Development install

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-projects

Cross-project example

python examples/quantum_chemistry/df_mp2_synthetic.py --demo

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

Status

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.

Roadmap

Five-phase plan from current alpha to generation-tuned stable release. See ROADMAP.md or the reader-oriented version on trnsci.dev.

Community

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.

About

Scientific computing library suite for AWS Trainium via NKI — the cuFFT/cuBLAS/cuRAND/cuSOLVER/cuSPARSE/cuTENSOR equivalents for Neuron. Python-first, PyTorch fallback everywhere, Apache-2.0.

Topics

Resources

License

Code of conduct

Contributing

Stars

Watchers

Forks

Packages

 
 
 

Contributors