You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Custom AWS Transform agent that migrates PyTorch/Triton kernels to AWS Trainium NKI (@nki.jit) and compiles, numerically verifies, and profiles every candidate on a real Trainium device before opening a PR.
FFT and complex-valued tensor operations for AWS Trainium via NKI (cuFFT-equivalent) — Cooley-Tukey, Bluestein, STFT, ComplexTensor and complex NN layers.
Random number generation for AWS Trainium via NKI (cuRAND-equivalent) — Philox counter-based RNG, standard distributions, Sobol / Halton / Latin-hypercube quasi-random sequences for Monte Carlo and QMC.
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.
Sparse matrix operations for AWS Trainium via NKI (cuSPARSE-equivalent) — CSR/COO formats, SpMV and SpMM via gather-matmul-scatter, Schwarz integral screening for quantum chemistry.
Tensor contractions for AWS Trainium via NKI (cuTENSOR-equivalent) — einsum with contraction planning, CP/PARAFAC and Tucker decompositions, density-fitted post-Hartree-Fock patterns.