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.
- trntensor-0.17.0/.github/workflows/ci.yml +59 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/.github/workflows/notify-umbrella.yml +3 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/.github/workflows/publish.yml +3 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/.gitignore +6 -0
- trntensor-0.17.0/.pre-commit-config.yaml +11 -0
- trntensor-0.17.0/CHANGELOG.md +686 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/CLAUDE.md +12 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/CONTRIBUTING.md +3 -2
- trntensor-0.17.0/PKG-INFO +144 -0
- trntensor-0.17.0/README.md +113 -0
- trntensor-0.17.0/benchmarks/bench_einsum.py +190 -0
- trntensor-0.17.0/docs/api/decompose.md +106 -0
- trntensor-0.17.0/docs/api/einsum.md +175 -0
- trntensor-0.17.0/docs/api/index.md +9 -0
- trntensor-0.17.0/docs/api/nki.md +78 -0
- trntensor-0.17.0/docs/api/plan.md +96 -0
- trntensor-0.17.0/docs/api/quantum.md +99 -0
- trntensor-0.17.0/docs/architecture.md +94 -0
- trntensor-0.17.0/docs/benchmarks.md +70 -0
- trntensor-0.17.0/docs/contraction_patterns.md +182 -0
- trntensor-0.17.0/docs/design/matmul_dispatch_findings.md +149 -0
- trntensor-0.17.0/docs/design/matmul_kernel_profile.md +159 -0
- trntensor-0.17.0/docs/developing_kernels.md +36 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/docs/quickstart.md +6 -6
- trntensor-0.17.0/docs/tutorial_df_mp2.md +157 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/examples/df_mp2_einsum.py +4 -3
- {trntensor-0.1.2 → trntensor-0.17.0}/infra/terraform/main.tf +2 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/mkdocs.yml +10 -1
- trntensor-0.17.0/pyproject.toml +92 -0
- trntensor-0.17.0/scripts/autotune_dispatch.py +233 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/scripts/bench_to_md.py +3 -8
- {trntensor-0.1.2 → trntensor-0.17.0}/scripts/run_benchmarks.sh +1 -1
- trntensor-0.17.0/scripts/run_neuron_profile.sh +469 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/scripts/run_neuron_tests.sh +17 -14
- trntensor-0.17.0/scripts/run_reduce_parallel_multirank.py +527 -0
- trntensor-0.17.0/scripts/run_simulator_tests.sh +137 -0
- trntensor-0.17.0/tests/conftest.py +43 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/tests/test_decompose.py +104 -9
- trntensor-0.17.0/tests/test_einsum.py +788 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/tests/test_nki.py +14 -0
- trntensor-0.17.0/tests/test_nki_compile.py +91 -0
- trntensor-0.17.0/tests/test_nki_kernels.py +424 -0
- trntensor-0.17.0/tests/test_nki_sim.py +264 -0
- trntensor-0.17.0/tests/test_parallel.py +451 -0
- trntensor-0.17.0/tests/test_plan.py +273 -0
- trntensor-0.17.0/tests/test_quantum.py +110 -0
- trntensor-0.17.0/tests/test_sr_runtime.py +163 -0
- trntensor-0.17.0/trntensor/__init__.py +59 -0
- trntensor-0.17.0/trntensor/decompose.py +278 -0
- trntensor-0.17.0/trntensor/einsum.py +637 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/trntensor/nki/__init__.py +3 -1
- trntensor-0.17.0/trntensor/nki/_kernels.py +362 -0
- trntensor-0.17.0/trntensor/nki/collectives.py +97 -0
- trntensor-0.17.0/trntensor/nki/dispatch.py +547 -0
- trntensor-0.17.0/trntensor/parallel.py +327 -0
- trntensor-0.17.0/trntensor/plan.py +486 -0
- trntensor-0.17.0/trntensor/py.typed +0 -0
- trntensor-0.17.0/trntensor/quantum.py +173 -0
- trntensor-0.17.0/trntensor.egg-info/PKG-INFO +144 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/trntensor.egg-info/SOURCES.txt +28 -1
- {trntensor-0.1.2 → trntensor-0.17.0}/trntensor.egg-info/requires.txt +3 -4
- trntensor-0.17.0/uv.lock +1637 -0
- trntensor-0.1.2/.github/workflows/ci.yml +0 -20
- trntensor-0.1.2/CHANGELOG.md +0 -56
- trntensor-0.1.2/PKG-INFO +0 -107
- trntensor-0.1.2/README.md +0 -76
- trntensor-0.1.2/benchmarks/bench_einsum.py +0 -56
- trntensor-0.1.2/docs/api.md +0 -27
- trntensor-0.1.2/docs/architecture.md +0 -34
- trntensor-0.1.2/docs/benchmarks.md +0 -38
- trntensor-0.1.2/pyproject.toml +0 -53
- trntensor-0.1.2/tests/conftest.py +0 -5
- trntensor-0.1.2/tests/test_einsum.py +0 -149
- trntensor-0.1.2/tests/test_plan.py +0 -111
- trntensor-0.1.2/trntensor/__init__.py +0 -21
- trntensor-0.1.2/trntensor/decompose.py +0 -180
- trntensor-0.1.2/trntensor/einsum.py +0 -110
- trntensor-0.1.2/trntensor/nki/dispatch.py +0 -40
- trntensor-0.1.2/trntensor/plan.py +0 -134
- trntensor-0.1.2/trntensor.egg-info/PKG-INFO +0 -107
- {trntensor-0.1.2 → trntensor-0.17.0}/CODE_OF_CONDUCT.md +0 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/LICENSE +0 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/benchmarks/__init__.py +0 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/docs/aws_setup.md +0 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/docs/index.md +0 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/docs/installation.md +0 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/infra/terraform/README.md +0 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/setup.cfg +0 -0
- {trntensor-0.1.2 → trntensor-0.17.0}/trntensor.egg-info/dependency_links.txt +0 -0
- {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
|
|
@@ -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
|