trntensor 0.3.0__tar.gz → 0.18.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 (86) hide show
  1. {trntensor-0.3.0 → trntensor-0.18.0}/.github/workflows/ci.yml +6 -4
  2. trntensor-0.18.0/CHANGELOG.md +777 -0
  3. {trntensor-0.3.0 → trntensor-0.18.0}/CLAUDE.md +4 -1
  4. {trntensor-0.3.0/trntensor.egg-info → trntensor-0.18.0}/PKG-INFO +40 -7
  5. {trntensor-0.3.0 → trntensor-0.18.0}/README.md +31 -5
  6. trntensor-0.18.0/docs/api/decompose.md +106 -0
  7. trntensor-0.18.0/docs/api/einsum.md +175 -0
  8. trntensor-0.18.0/docs/api/nki.md +174 -0
  9. trntensor-0.18.0/docs/api/plan.md +96 -0
  10. {trntensor-0.3.0 → trntensor-0.18.0}/docs/api/quantum.md +16 -6
  11. {trntensor-0.3.0 → trntensor-0.18.0}/docs/architecture.md +34 -3
  12. {trntensor-0.3.0 → trntensor-0.18.0}/docs/contraction_patterns.md +19 -18
  13. trntensor-0.18.0/docs/design/matmul_dispatch_findings.md +149 -0
  14. trntensor-0.18.0/docs/design/matmul_kernel_profile.md +159 -0
  15. {trntensor-0.3.0 → trntensor-0.18.0}/docs/quickstart.md +6 -6
  16. {trntensor-0.3.0 → trntensor-0.18.0}/docs/tutorial_df_mp2.md +2 -4
  17. {trntensor-0.3.0 → trntensor-0.18.0}/infra/terraform/main.tf +3 -1
  18. {trntensor-0.3.0 → trntensor-0.18.0}/pyproject.toml +29 -8
  19. trntensor-0.18.0/scripts/autotune_dispatch.py +383 -0
  20. trntensor-0.18.0/scripts/bench_native_einsum.py +396 -0
  21. trntensor-0.18.0/scripts/run_neuron_profile.sh +469 -0
  22. trntensor-0.18.0/scripts/run_reduce_parallel_multirank.py +527 -0
  23. trntensor-0.18.0/tests/conftest.py +49 -0
  24. trntensor-0.18.0/tests/test_einsum.py +788 -0
  25. trntensor-0.18.0/tests/test_native_dispatch.py +732 -0
  26. trntensor-0.18.0/tests/test_nki_compile.py +91 -0
  27. {trntensor-0.3.0 → trntensor-0.18.0}/tests/test_nki_kernels.py +128 -52
  28. trntensor-0.18.0/tests/test_nki_native.py +176 -0
  29. {trntensor-0.3.0 → trntensor-0.18.0}/tests/test_nki_sim.py +111 -1
  30. trntensor-0.18.0/tests/test_parallel.py +451 -0
  31. {trntensor-0.3.0 → trntensor-0.18.0}/tests/test_plan.py +44 -2
  32. {trntensor-0.3.0 → trntensor-0.18.0}/tests/test_quantum.py +24 -0
  33. trntensor-0.18.0/tests/test_sr_runtime.py +166 -0
  34. {trntensor-0.3.0 → trntensor-0.18.0}/trntensor/__init__.py +12 -2
  35. trntensor-0.18.0/trntensor/einsum.py +649 -0
  36. trntensor-0.18.0/trntensor/nki/_kernels.py +362 -0
  37. trntensor-0.18.0/trntensor/nki/_native.py +85 -0
  38. trntensor-0.18.0/trntensor/nki/collectives.py +97 -0
  39. trntensor-0.18.0/trntensor/nki/dispatch.py +911 -0
  40. trntensor-0.18.0/trntensor/parallel.py +327 -0
  41. {trntensor-0.3.0 → trntensor-0.18.0}/trntensor/plan.py +195 -14
  42. {trntensor-0.3.0 → trntensor-0.18.0}/trntensor/quantum.py +52 -45
  43. {trntensor-0.3.0 → trntensor-0.18.0/trntensor.egg-info}/PKG-INFO +40 -7
  44. {trntensor-0.3.0 → trntensor-0.18.0}/trntensor.egg-info/SOURCES.txt +14 -0
  45. {trntensor-0.3.0 → trntensor-0.18.0}/trntensor.egg-info/requires.txt +10 -1
  46. {trntensor-0.3.0 → trntensor-0.18.0}/uv.lock +1 -1
  47. trntensor-0.3.0/CHANGELOG.md +0 -201
  48. trntensor-0.3.0/docs/api/decompose.md +0 -41
  49. trntensor-0.3.0/docs/api/einsum.md +0 -43
  50. trntensor-0.3.0/docs/api/nki.md +0 -78
  51. trntensor-0.3.0/docs/api/plan.md +0 -36
  52. trntensor-0.3.0/tests/conftest.py +0 -12
  53. trntensor-0.3.0/tests/test_einsum.py +0 -293
  54. trntensor-0.3.0/trntensor/einsum.py +0 -228
  55. trntensor-0.3.0/trntensor/nki/_kernels.py +0 -296
  56. trntensor-0.3.0/trntensor/nki/dispatch.py +0 -343
  57. {trntensor-0.3.0 → trntensor-0.18.0}/.github/workflows/notify-umbrella.yml +0 -0
  58. {trntensor-0.3.0 → trntensor-0.18.0}/.github/workflows/publish.yml +0 -0
  59. {trntensor-0.3.0 → trntensor-0.18.0}/.gitignore +0 -0
  60. {trntensor-0.3.0 → trntensor-0.18.0}/.pre-commit-config.yaml +0 -0
  61. {trntensor-0.3.0 → trntensor-0.18.0}/CODE_OF_CONDUCT.md +0 -0
  62. {trntensor-0.3.0 → trntensor-0.18.0}/CONTRIBUTING.md +0 -0
  63. {trntensor-0.3.0 → trntensor-0.18.0}/LICENSE +0 -0
  64. {trntensor-0.3.0 → trntensor-0.18.0}/benchmarks/__init__.py +0 -0
  65. {trntensor-0.3.0 → trntensor-0.18.0}/benchmarks/bench_einsum.py +0 -0
  66. {trntensor-0.3.0 → trntensor-0.18.0}/docs/api/index.md +0 -0
  67. {trntensor-0.3.0 → trntensor-0.18.0}/docs/aws_setup.md +0 -0
  68. {trntensor-0.3.0 → trntensor-0.18.0}/docs/benchmarks.md +0 -0
  69. {trntensor-0.3.0 → trntensor-0.18.0}/docs/developing_kernels.md +0 -0
  70. {trntensor-0.3.0 → trntensor-0.18.0}/docs/index.md +0 -0
  71. {trntensor-0.3.0 → trntensor-0.18.0}/docs/installation.md +0 -0
  72. {trntensor-0.3.0 → trntensor-0.18.0}/examples/df_mp2_einsum.py +0 -0
  73. {trntensor-0.3.0 → trntensor-0.18.0}/infra/terraform/README.md +0 -0
  74. {trntensor-0.3.0 → trntensor-0.18.0}/mkdocs.yml +0 -0
  75. {trntensor-0.3.0 → trntensor-0.18.0}/scripts/bench_to_md.py +0 -0
  76. {trntensor-0.3.0 → trntensor-0.18.0}/scripts/run_benchmarks.sh +0 -0
  77. {trntensor-0.3.0 → trntensor-0.18.0}/scripts/run_neuron_tests.sh +0 -0
  78. {trntensor-0.3.0 → trntensor-0.18.0}/scripts/run_simulator_tests.sh +0 -0
  79. {trntensor-0.3.0 → trntensor-0.18.0}/setup.cfg +0 -0
  80. {trntensor-0.3.0 → trntensor-0.18.0}/tests/test_decompose.py +0 -0
  81. {trntensor-0.3.0 → trntensor-0.18.0}/tests/test_nki.py +0 -0
  82. {trntensor-0.3.0 → trntensor-0.18.0}/trntensor/decompose.py +0 -0
  83. {trntensor-0.3.0 → trntensor-0.18.0}/trntensor/nki/__init__.py +0 -0
  84. {trntensor-0.3.0 → trntensor-0.18.0}/trntensor/py.typed +0 -0
  85. {trntensor-0.3.0 → trntensor-0.18.0}/trntensor.egg-info/dependency_links.txt +0 -0
  86. {trntensor-0.3.0 → trntensor-0.18.0}/trntensor.egg-info/top_level.txt +0 -0
{trntensor-0.3.0 → trntensor-0.18.0}/.github/workflows/ci.yml RENAMED
@@ -30,8 +30,10 @@ jobs:
30
30
  - uses: astral-sh/setup-uv@v6
31
31
  with:
32
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
33
+ # uv pip (not uv sync): sync resolves every extra, and [torchneuron]'s nki>=0.6
34
+ # lives only on the Neuron index.
35
+ - run: uv venv && uv pip install -e ".[dev]"
36
+ - run: uv run --no-sync pytest tests/ -v -x --tb=short -m "not neuron and not nki_simulator" --cov=trntensor --cov-report=xml
35
37
 
36
38
  - name: Upload coverage reports to Codecov
37
39
  uses: codecov/codecov-action@v5
@@ -50,10 +52,10 @@ jobs:
50
52
  - uses: astral-sh/setup-uv@v6
51
53
  with:
52
54
  python-version: "3.12"
53
- - run: uv sync --extra dev
55
+ - run: uv venv && uv pip install -e ".[dev]"
54
56
  - name: Install NKI simulator
55
57
  run: uv pip install --extra-index-url https://pip.repos.neuron.amazonaws.com "nki>=0.3.0"
56
58
  - name: Run simulator-backed kernel tests
57
59
  env:
58
60
  TRNTENSOR_USE_SIMULATOR: "1"
59
- run: uv run pytest tests/ -v -m nki_simulator --tb=short
61
+ run: uv run --no-sync pytest tests/ -v -m nki_simulator --tb=short