trntensor 0.1.2__tar.gz → 0.17.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (90) hide show
  1. trntensor-0.17.0/.github/workflows/ci.yml +59 -0
  2. {trntensor-0.1.2 → trntensor-0.17.0}/.github/workflows/notify-umbrella.yml +3 -0
  3. {trntensor-0.1.2 → trntensor-0.17.0}/.github/workflows/publish.yml +3 -0
  4. {trntensor-0.1.2 → trntensor-0.17.0}/.gitignore +6 -0
  5. trntensor-0.17.0/.pre-commit-config.yaml +11 -0
  6. trntensor-0.17.0/CHANGELOG.md +686 -0
  7. {trntensor-0.1.2 → trntensor-0.17.0}/CLAUDE.md +12 -0
  8. {trntensor-0.1.2 → trntensor-0.17.0}/CONTRIBUTING.md +3 -2
  9. trntensor-0.17.0/PKG-INFO +144 -0
  10. trntensor-0.17.0/README.md +113 -0
  11. trntensor-0.17.0/benchmarks/bench_einsum.py +190 -0
  12. trntensor-0.17.0/docs/api/decompose.md +106 -0
  13. trntensor-0.17.0/docs/api/einsum.md +175 -0
  14. trntensor-0.17.0/docs/api/index.md +9 -0
  15. trntensor-0.17.0/docs/api/nki.md +78 -0
  16. trntensor-0.17.0/docs/api/plan.md +96 -0
  17. trntensor-0.17.0/docs/api/quantum.md +99 -0
  18. trntensor-0.17.0/docs/architecture.md +94 -0
  19. trntensor-0.17.0/docs/benchmarks.md +70 -0
  20. trntensor-0.17.0/docs/contraction_patterns.md +182 -0
  21. trntensor-0.17.0/docs/design/matmul_dispatch_findings.md +149 -0
  22. trntensor-0.17.0/docs/design/matmul_kernel_profile.md +159 -0
  23. trntensor-0.17.0/docs/developing_kernels.md +36 -0
  24. {trntensor-0.1.2 → trntensor-0.17.0}/docs/quickstart.md +6 -6
  25. trntensor-0.17.0/docs/tutorial_df_mp2.md +157 -0
  26. {trntensor-0.1.2 → trntensor-0.17.0}/examples/df_mp2_einsum.py +4 -3
  27. {trntensor-0.1.2 → trntensor-0.17.0}/infra/terraform/main.tf +2 -0
  28. {trntensor-0.1.2 → trntensor-0.17.0}/mkdocs.yml +10 -1
  29. trntensor-0.17.0/pyproject.toml +92 -0
  30. trntensor-0.17.0/scripts/autotune_dispatch.py +233 -0
  31. {trntensor-0.1.2 → trntensor-0.17.0}/scripts/bench_to_md.py +3 -8
  32. {trntensor-0.1.2 → trntensor-0.17.0}/scripts/run_benchmarks.sh +1 -1
  33. trntensor-0.17.0/scripts/run_neuron_profile.sh +469 -0
  34. {trntensor-0.1.2 → trntensor-0.17.0}/scripts/run_neuron_tests.sh +17 -14
  35. trntensor-0.17.0/scripts/run_reduce_parallel_multirank.py +527 -0
  36. trntensor-0.17.0/scripts/run_simulator_tests.sh +137 -0
  37. trntensor-0.17.0/tests/conftest.py +43 -0
  38. {trntensor-0.1.2 → trntensor-0.17.0}/tests/test_decompose.py +104 -9
  39. trntensor-0.17.0/tests/test_einsum.py +788 -0
  40. {trntensor-0.1.2 → trntensor-0.17.0}/tests/test_nki.py +14 -0
  41. trntensor-0.17.0/tests/test_nki_compile.py +91 -0
  42. trntensor-0.17.0/tests/test_nki_kernels.py +424 -0
  43. trntensor-0.17.0/tests/test_nki_sim.py +264 -0
  44. trntensor-0.17.0/tests/test_parallel.py +451 -0
  45. trntensor-0.17.0/tests/test_plan.py +273 -0
  46. trntensor-0.17.0/tests/test_quantum.py +110 -0
  47. trntensor-0.17.0/tests/test_sr_runtime.py +163 -0
  48. trntensor-0.17.0/trntensor/__init__.py +59 -0
  49. trntensor-0.17.0/trntensor/decompose.py +278 -0
  50. trntensor-0.17.0/trntensor/einsum.py +637 -0
  51. {trntensor-0.1.2 → trntensor-0.17.0}/trntensor/nki/__init__.py +3 -1
  52. trntensor-0.17.0/trntensor/nki/_kernels.py +362 -0
  53. trntensor-0.17.0/trntensor/nki/collectives.py +97 -0
  54. trntensor-0.17.0/trntensor/nki/dispatch.py +547 -0
  55. trntensor-0.17.0/trntensor/parallel.py +327 -0
  56. trntensor-0.17.0/trntensor/plan.py +486 -0
  57. trntensor-0.17.0/trntensor/py.typed +0 -0
  58. trntensor-0.17.0/trntensor/quantum.py +173 -0
  59. trntensor-0.17.0/trntensor.egg-info/PKG-INFO +144 -0
  60. {trntensor-0.1.2 → trntensor-0.17.0}/trntensor.egg-info/SOURCES.txt +28 -1
  61. {trntensor-0.1.2 → trntensor-0.17.0}/trntensor.egg-info/requires.txt +3 -4
  62. trntensor-0.17.0/uv.lock +1637 -0
  63. trntensor-0.1.2/.github/workflows/ci.yml +0 -20
  64. trntensor-0.1.2/CHANGELOG.md +0 -56
  65. trntensor-0.1.2/PKG-INFO +0 -107
  66. trntensor-0.1.2/README.md +0 -76
  67. trntensor-0.1.2/benchmarks/bench_einsum.py +0 -56
  68. trntensor-0.1.2/docs/api.md +0 -27
  69. trntensor-0.1.2/docs/architecture.md +0 -34
  70. trntensor-0.1.2/docs/benchmarks.md +0 -38
  71. trntensor-0.1.2/pyproject.toml +0 -53
  72. trntensor-0.1.2/tests/conftest.py +0 -5
  73. trntensor-0.1.2/tests/test_einsum.py +0 -149
  74. trntensor-0.1.2/tests/test_plan.py +0 -111
  75. trntensor-0.1.2/trntensor/__init__.py +0 -21
  76. trntensor-0.1.2/trntensor/decompose.py +0 -180
  77. trntensor-0.1.2/trntensor/einsum.py +0 -110
  78. trntensor-0.1.2/trntensor/nki/dispatch.py +0 -40
  79. trntensor-0.1.2/trntensor/plan.py +0 -134
  80. trntensor-0.1.2/trntensor.egg-info/PKG-INFO +0 -107
  81. {trntensor-0.1.2 → trntensor-0.17.0}/CODE_OF_CONDUCT.md +0 -0
  82. {trntensor-0.1.2 → trntensor-0.17.0}/LICENSE +0 -0
  83. {trntensor-0.1.2 → trntensor-0.17.0}/benchmarks/__init__.py +0 -0
  84. {trntensor-0.1.2 → trntensor-0.17.0}/docs/aws_setup.md +0 -0
  85. {trntensor-0.1.2 → trntensor-0.17.0}/docs/index.md +0 -0
  86. {trntensor-0.1.2 → trntensor-0.17.0}/docs/installation.md +0 -0
  87. {trntensor-0.1.2 → trntensor-0.17.0}/infra/terraform/README.md +0 -0
  88. {trntensor-0.1.2 → trntensor-0.17.0}/setup.cfg +0 -0
  89. {trntensor-0.1.2 → trntensor-0.17.0}/trntensor.egg-info/dependency_links.txt +0 -0
  90. {trntensor-0.1.2 → trntensor-0.17.0}/trntensor.egg-info/top_level.txt +0 -0
@@ -0,0 +1,59 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+ pull_request:
7
+
8
+ # Opt into Node.js 24 for all third-party JavaScript actions ahead of
9
+ # GitHub's June 2026 default switch. Silences the Node 20 deprecation
10
+ # notice emitted by setup-uv and codecov-action's github-script dep.
11
+ env:
12
+ FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true
13
+
14
+ jobs:
15
+ lint:
16
+ runs-on: ubuntu-latest
17
+ steps:
18
+ - uses: actions/checkout@v6
19
+ - uses: astral-sh/setup-uv@v6
20
+ - run: uvx ruff check .
21
+ - run: uvx ruff format --check .
22
+
23
+ test:
24
+ runs-on: ubuntu-latest
25
+ strategy:
26
+ matrix:
27
+ python-version: ["3.10", "3.11", "3.12"]
28
+ steps:
29
+ - uses: actions/checkout@v6
30
+ - uses: astral-sh/setup-uv@v6
31
+ with:
32
+ python-version: ${{ matrix.python-version }}
33
+ - run: uv sync --extra dev
34
+ - run: uv run pytest tests/ -v -x --tb=short -m "not neuron and not nki_simulator" --cov=trntensor --cov-report=xml
35
+
36
+ - name: Upload coverage reports to Codecov
37
+ uses: codecov/codecov-action@v5
38
+ with:
39
+ token: ${{ secrets.CODECOV_TOKEN }}
40
+ slug: trnsci/trntensor
41
+
42
+ nki-simulator:
43
+ # Runs NKI kernels through nki.simulate(kernel)(numpy_args) on CPU.
44
+ # Catches Python-trace-level errors (bad kwargs, dropped ops, shape
45
+ # mismatches) pre-merge without AWS round-trips. MLIR verifier
46
+ # errors remain hardware-only (simulator explicitly skips compile).
47
+ runs-on: ubuntu-latest
48
+ steps:
49
+ - uses: actions/checkout@v6
50
+ - uses: astral-sh/setup-uv@v6
51
+ with:
52
+ python-version: "3.12"
53
+ - run: uv sync --extra dev
54
+ - name: Install NKI simulator
55
+ run: uv pip install --extra-index-url https://pip.repos.neuron.amazonaws.com "nki>=0.3.0"
56
+ - name: Run simulator-backed kernel tests
57
+ env:
58
+ TRNTENSOR_USE_SIMULATOR: "1"
59
+ run: uv run pytest tests/ -v -m nki_simulator --tb=short
@@ -18,6 +18,9 @@ on:
18
18
  - "CHANGELOG.md"
19
19
  workflow_dispatch:
20
20
 
21
+ env:
22
+ FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true
23
+
21
24
  jobs:
22
25
  dispatch:
23
26
  runs-on: ubuntu-latest
@@ -4,6 +4,9 @@ on:
4
4
  release:
5
5
  types: [published]
6
6
 
7
+ env:
8
+ FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true
9
+
7
10
  jobs:
8
11
  build:
9
12
  runs-on: ubuntu-latest
@@ -31,3 +31,9 @@ site/
31
31
  # coverage artifacts
32
32
  .coverage
33
33
  htmlcov/
34
+
35
+ # terraform local state
36
+ infra/terraform/.terraform/
37
+ infra/terraform/terraform.tfstate*
38
+ infra/terraform/.terraform.lock.hcl
39
+ benchmark_results.json
@@ -0,0 +1,11 @@
1
+ # Install locally: pip install pre-commit && pre-commit install
2
+ # Runs on `git commit`; ruff auto-fixes staged Python files then formats them.
3
+ # CI re-checks the same rules in the lint job — pre-commit keeps CI green.
4
+
5
+ repos:
6
+ - repo: https://github.com/astral-sh/ruff-pre-commit
7
+ rev: v0.6.9
8
+ hooks:
9
+ - id: ruff
10
+ args: [--fix]
11
+ - id: ruff-format