{"id":48762154,"url":"https://github.com/trnsci/trntensor","last_synced_at":"2026-04-19T01:01:03.082Z","repository":{"id":350951974,"uuid":"1208919744","full_name":"trnsci/trntensor","owner":"trnsci","description":"Tensor contractions for AWS Trainium via NKI (cuTENSOR-equivalent) — einsum with contraction planning, CP/PARAFAC and Tucker decompositions, density-fitted post-Hartree-Fock patterns.","archived":false,"fork":false,"pushed_at":"2026-04-17T00:08:36.000Z","size":309,"stargazers_count":0,"open_issues_count":3,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2026-04-18T00:03:50.082Z","etag":null,"topics":["aws-neuron","aws-trainium","cp-decomposition","cutensor","einsum","nki","python","pytorch","scientific-computing","tensor-contraction","tensor-decomposition","tucker-decomposition"],"latest_commit_sha":null,"homepage":"https://trnsci.dev/trntensor/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"other","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/trnsci.png","metadata":{"files":{"readme":"README.md","changelog":"CHANGELOG.md","contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":"CODE_OF_CONDUCT.md","threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2026-04-12T23:03:11.000Z","updated_at":"2026-04-17T00:08:40.000Z","dependencies_parsed_at":"2026-04-18T00:01:02.330Z","dependency_job_id":null,"html_url":"https://github.com/trnsci/trntensor","commit_stats":null,"previous_names":["trnsci/trntensor"],"tags_count":5,"template":false,"template_full_name":null,"purl":"pkg:github/trnsci/trntensor","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/trnsci%2Ftrntensor","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/trnsci%2Ftrntensor/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/trnsci%2Ftrntensor/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/trnsci%2Ftrntensor/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/trnsci","download_url":"https://codeload.github.com/trnsci/trntensor/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/trnsci%2Ftrntensor/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":31990577,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-04-18T20:23:30.271Z","status":"ssl_error","status_checked_at":"2026-04-18T20:23:29.375Z","response_time":103,"last_error":"SSL_read: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["aws-neuron","aws-trainium","cp-decomposition","cutensor","einsum","nki","python","pytorch","scientific-computing","tensor-contraction","tensor-decomposition","tucker-decomposition"],"created_at":"2026-04-13T07:38:57.485Z","updated_at":"2026-04-19T01:01:03.040Z","avatar_url":"https://github.com/trnsci.png","language":"Python","readme":"# trntensor\n\n[![CI](https://github.com/trnsci/trntensor/actions/workflows/ci.yml/badge.svg)](https://github.com/trnsci/trntensor/actions/workflows/ci.yml)\n[![codecov](https://codecov.io/gh/trnsci/trntensor/graph/badge.svg)](https://codecov.io/gh/trnsci/trntensor)\n[![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)\n[![PyPI](https://img.shields.io/pypi/v/trntensor)](https://pypi.org/project/trntensor/)\n[![Python](https://img.shields.io/pypi/pyversions/trntensor)](https://pypi.org/project/trntensor/)\n[![License](https://img.shields.io/github/license/trnsci/trntensor)](LICENSE)\n[![Docs](https://img.shields.io/badge/docs-trnsci.dev-blue)](https://trnsci.dev/trntensor/)\n\nTensor contractions for AWS Trainium via NKI.\n\nEinstein 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](https://github.com/trnsci)).\n\n## Current phase\n\ntrntensor follows the [trnsci 5-phase roadmap](https://trnsci.dev/roadmap/). Active work is tracked in phase-labeled GitHub issues:\n\n- **[Phase 1 — correctness](https://github.com/trnsci/trntensor/issues/27)** (active): matmul + batched-matmul NKI kernels in place; awaiting hardware validation + additional `@pytest.mark.neuron` coverage.\n- **[Phase 2 — precision](https://github.com/trnsci/trntensor/issues/28)**: precision-aware contraction path selection (depends on [trnblas#22](https://github.com/trnsci/trnblas/issues/22) double-double GEMM).\n- **[Phase 3 — perf](https://github.com/trnsci/trntensor/issues/29)**: opt_einsum-style path planner, plan cache reuse.\n- **[Phase 4 — multi-chip](https://github.com/trnsci/trntensor/issues/30)**: sharded tensor contractions.\n- **[Phase 5 — generation](https://github.com/trnsci/trntensor/issues/31)**: trn2 fused multi-contraction paths.\n\nSuite-wide tracker: [trnsci/trnsci#1](https://github.com/trnsci/trnsci/issues/1).\n\n## Install\n\n```bash\npip install trntensor\n```\n\n## Usage\n\n```python\nimport torch\nimport trntensor\n\n# Einsum — drop-in for torch.einsum with contraction planning\nC = trntensor.einsum(\"ij,jk-\u003eik\", A, B)           # matmul\nT = trntensor.einsum(\"ap,bp-\u003eab\", B_i, B_j)       # DF-MP2 contraction\nX = trntensor.einsum(\"mi,mnP-\u003einP\", C_occ, eri)   # AO→MO transform\n\n# Contraction planning\nplan = trntensor.plan_contraction(\"ij,jk-\u003eik\", A, B)\nflops = trntensor.estimate_flops(\"ij,jk-\u003eik\", A, B)\n\n# CP decomposition (tensor hypercontraction)\nfactors, weights = trntensor.cp_decompose(tensor, rank=10)\nreconstructed = trntensor.cp_reconstruct(factors, weights)\n\n# Tucker decomposition (HOSVD)\ncore, factors = trntensor.tucker_decompose(tensor, ranks=(5, 5, 5))\n```\n\n## Operations\n\n| Category | Operation | Description |\n|----------|-----------|-------------|\n| Contraction | `einsum` | General tensor contraction |\n| Contraction | `multi_einsum` | Multiple contractions (fusion-ready) |\n| Planning | `plan_contraction` | Analyze and select strategy |\n| Planning | `estimate_flops` | FLOPs for a contraction |\n| Decomposition | `cp_decompose` | CP/PARAFAC via ALS |\n| Decomposition | `tucker_decompose` | Tucker via HOSVD |\n\n## Status\n\n- [x] Einsum with matmul/bmm/torch dispatch\n- [x] Contraction planner\n- [x] CP decomposition (ALS)\n- [x] Tucker decomposition (HOSVD)\n- [x] DF-MP2 einsum example\n- [x] NKI fused contraction kernels (`mp2_energy`, `ao_to_mo_transform`)\n- [x] XLA operand residency (`to_xla` / `from_xla`)\n- [x] NKI CPU simulator + `nki-simulator` CI gate\n- [x] Optimal contraction ordering — greedy path search for 3+ operands\n- [x] Multi-contraction shared-operand XLA pre-pinning\n- [x] Contraction plan cache (`clear_plan_cache` / `plan_cache_info`)\n- [x] Alpha/beta scaling for `einsum` (cuTENSOR GEMM-style)\n- [x] Input validation with descriptive errors\n- [x] Tensor Train (TT) decomposition (TT-SVD)\n- [x] Non-negative CP + warm-start CP\n- [x] PEP 561 `py.typed` marker\n- [x] Mixed precision / `dtype` override for bf16/fp16\n- [x] `precision=` kwarg (`\"fast\"` / `\"kahan\"` fp64 promotion / `\"dd\"` stub)\n\n## Related Projects\n\n| Project | What |\n|---------|------|\n| [trnfft](https://github.com/scttfrdmn/trnfft) | FFT + complex ops |\n| [trnblas](https://github.com/scttfrdmn/trnblas) | BLAS operations |\n| [trnsolver](https://github.com/scttfrdmn/trnsolver) | Linear solvers |\n| [trnrand](https://github.com/scttfrdmn/trnrand) | Random number generation |\n| [trnsparse](https://github.com/scttfrdmn/trnsparse) | Sparse operations |\n\n## License\n\nApache 2.0 — Copyright 2026 Scott Friedman\n\n\n## Disclaimer\n\ntrnsci 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.\n\n\"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.\n\nAll work, opinions, analyses, benchmark results, architectural commentary, and editorial judgments in this repository and on [trnsci.dev](https://trnsci.dev) are those of the project's contributors. They do not represent the views, positions, or commitments of Amazon, AWS, or Annapurna Labs.\n\nFeedback 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.\n\nFor official AWS guidance, see [aws-neuron documentation](https://awsdocs-neuron.readthedocs-hosted.com/) and the [AWS Trainium product page](https://aws.amazon.com/ai/machine-learning/trainium/).\n","funding_links":[],"categories":[],"sub_categories":[],"project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftrnsci%2Ftrntensor","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ftrnsci%2Ftrntensor","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftrnsci%2Ftrntensor/lists"}