haliax 1.4.dev366__tar.gz → 1.4.dev367__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.
- {haliax-1.4.dev366 → haliax-1.4.dev367}/AGENTS.md +3 -9
- {haliax-1.4.dev366 → haliax-1.4.dev367}/PKG-INFO +1 -15
- {haliax-1.4.dev366 → haliax-1.4.dev367}/pyproject.toml +11 -4
- haliax-1.4.dev367/src/haliax/__about__.py +1 -0
- haliax-1.4.dev366/src/haliax/__about__.py +0 -1
- {haliax-1.4.dev366 → haliax-1.4.dev367}/.coveragerc +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/.flake8 +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/.github/workflows/publish_dev.yaml +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/.github/workflows/run_pre_commit.yaml +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/.github/workflows/run_quick_levanter_tests.yaml +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/.github/workflows/run_tests.yaml +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/.gitignore +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/.pre-commit-config.yaml +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/.readthedocs.yaml +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/CONTRIBUTING.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/LICENSE +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/README.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/api.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/broadcasting.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/cheatsheet.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/css/material.css +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/css/mkdocstrings.css +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/faq.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/data_parallel_mesh.png +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/data_parallel_mesh_replicated.png +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_1d.png +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_1d_zero.png +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_2d.png +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_2d_batch_partitioned.png +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_2d_data_replicated.png +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_2d_data_replicated_mlp_partitioned.png +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_2d_intermediate_fully_partitioned.png +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_2d_zero.png +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/fp8.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/index.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/indexing.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/matmul.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/nn.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/partitioning.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/rearrange.ipynb +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/rearrange.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/requirements.txt +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/scan.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/state-dict.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/tutorial.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/typing.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/vmap.md +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/mkdocs.yml +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/__init__.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/__init__.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/compile_utils.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/dot.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/einsum.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/fp8.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/parsing.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/rearrange.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/scan.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/state_dict.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/util.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/axis.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/core.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/debug.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/haxtyping.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/hof.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/jax_utils.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/__init__.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/activations.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/attention.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/conv.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/dropout.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/embedding.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/linear.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/loss.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/mlp.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/normalization.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/pool.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/scan.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/ops.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/partitioning.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/quantization.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/random.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/specialized_fns.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/state_dict.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/tree_util.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/types.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/util.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/wrap.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/core_test.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_attention.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_axis.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_conv.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_debug.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_dot.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_dtype_typing.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_einsum.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_fp8.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_hof.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_int8.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_namedarray_typing.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_nn.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_ops.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_parsing.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_partitioning.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_pool.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_random.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_rearrange.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_scan.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_scatter_gather.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_specialized_fns.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_state_dict.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_tree_util.py +0 -0
- {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_utils.py +0 -0
|
@@ -25,7 +25,8 @@ repository. Follow these notes when implementing new features or fixing bugs.
|
|
|
25
25
|
* **Formatting and Linting**: We use `ruff` via `pre-commit`.
|
|
26
26
|
* **Typing**: the code base uses `mypy` for static type checking. `mypy` is run by pre‑commit and the
|
|
27
27
|
configuration is found in `pyproject.toml`.
|
|
28
|
-
* **Run `pre-commit run --all-files`** before committing. The CI workflows run the same checks.
|
|
28
|
+
* **Run `uv run pre-commit run --all-files`** before committing. The CI workflows run the same checks.
|
|
29
|
+
* **Use `uv run` for commands.** When running tools like `pytest` or other scripts, invoke them via `uv run` so the development dependencies are active.
|
|
29
30
|
* **Doc Strings**: All public functions, classes, and modules should have docstrings, unless
|
|
30
31
|
their purpose is painfully obvious. Use
|
|
31
32
|
[Google style](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) for
|
|
@@ -44,12 +45,10 @@ repository. Follow these notes when implementing new features or fixing bugs.
|
|
|
44
45
|
|
|
45
46
|
## Testing
|
|
46
47
|
|
|
47
|
-
* Tests are executed with `pytest`. The default workflow runs
|
|
48
|
-
`pytest tests -m "not entry and not slow and not ray"`.
|
|
48
|
+
* Tests are executed with `pytest`. The default workflow runs `uv run pytest tests`.
|
|
49
49
|
* In general, never relax tolerances in floating point tests unless specifically discussed with the
|
|
50
50
|
team. Use `assert_allclose` with appropriate tolerances for numerical comparisons. We typically use
|
|
51
51
|
1e-4 for more complex modules, and 1e-5 for simpler ones.
|
|
52
|
-
* Tests should be reasonably fast. Mark long-running tests with @pytest.mark.slow so they are excluded from the default suite.
|
|
53
52
|
* Always mark tests that depend on pytorch with `@skip_if_no_torch` to ensure they are skipped
|
|
54
53
|
when PyTorch is not available. This is particularly important for tests that require PyTorch-specific
|
|
55
54
|
functionality.
|
|
@@ -59,11 +58,6 @@ repository. Follow these notes when implementing new features or fixing bugs.
|
|
|
59
58
|
|
|
60
59
|
* **Generic code**: many utilities are written with Python generics and dataclasses. Where possible,
|
|
61
60
|
write reusable functions or classes that operate over TypeVars instead of hard coding concrete types.
|
|
62
|
-
* **Configurations**: configuration files are dataclasses loaded via `draccus`. Keep configs
|
|
63
|
-
declarative and typed.
|
|
64
|
-
* **Reproducibility**: Levanter aims for deterministic training where possible. Avoid sources of
|
|
65
|
-
nondeterminism unless explicitly required.
|
|
66
|
-
* Prefer Stacked with fold or scan over writing custom loops, for better compile times and gradient checkpointing support
|
|
67
61
|
|
|
68
62
|
## Library conventions
|
|
69
63
|
- Haliax revolves around `NamedArray` and explicit `Axis` objects. Prefer APIs that accept
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: haliax
|
|
3
|
-
Version: 1.4.
|
|
3
|
+
Version: 1.4.dev367
|
|
4
4
|
Summary: Named Tensors for Legible Deep Learning in JAX
|
|
5
5
|
Project-URL: Homepage, https://github.com/stanford-crfm/haliax
|
|
6
6
|
Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/
|
|
@@ -19,20 +19,6 @@ Requires-Dist: equinox>=0.10.6
|
|
|
19
19
|
Requires-Dist: jaxtyping>=0.2.20
|
|
20
20
|
Requires-Dist: jmp>=0.0.4
|
|
21
21
|
Requires-Dist: safetensors>=0.4.3
|
|
22
|
-
Provides-Extra: dev
|
|
23
|
-
Requires-Dist: chex>=0.1.86; extra == 'dev'
|
|
24
|
-
Requires-Dist: mkdocs-include-markdown-plugin; extra == 'dev'
|
|
25
|
-
Requires-Dist: mkdocs-literate-nav>=0.6.0; extra == 'dev'
|
|
26
|
-
Requires-Dist: mkdocs-macros-plugin>=0.7.0; extra == 'dev'
|
|
27
|
-
Requires-Dist: mkdocs-material>=7.3.3; extra == 'dev'
|
|
28
|
-
Requires-Dist: mkdocs>=1.4.3; extra == 'dev'
|
|
29
|
-
Requires-Dist: mkdocstrings-python>=1.1.2; extra == 'dev'
|
|
30
|
-
Requires-Dist: mkdocstrings>=0.22.0; extra == 'dev'
|
|
31
|
-
Requires-Dist: mypy>=0.910; extra == 'dev'
|
|
32
|
-
Requires-Dist: pre-commit; extra == 'dev'
|
|
33
|
-
Requires-Dist: pygments; extra == 'dev'
|
|
34
|
-
Requires-Dist: pymdown-extensions; extra == 'dev'
|
|
35
|
-
Requires-Dist: pytest>=7.4.0; extra == 'dev'
|
|
36
22
|
Description-Content-Type: text/markdown
|
|
37
23
|
|
|
38
24
|
<!--haliax-intro-start-->
|
|
@@ -31,13 +31,20 @@ dependencies = [
|
|
|
31
31
|
]
|
|
32
32
|
dynamic =[ "version" ]
|
|
33
33
|
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
34
|
+
|
|
35
|
+
[dependency-groups]
|
|
36
|
+
dev = [
|
|
37
|
+
"pytest >= 7.4.0",
|
|
38
|
+
"mypy >= 0.910",
|
|
39
|
+
"mkdocs >= 1.4.3",
|
|
40
|
+
"mkdocs-material >= 7.3.3",
|
|
41
|
+
"mkdocstrings >= 0.22.0",
|
|
42
|
+
"mkdocs-literate-nav >= 0.6.0",
|
|
43
|
+
"mkdocs-macros-plugin >= 0.7.0",
|
|
44
|
+
"mkdocstrings-python >= 1.1.2",
|
|
37
45
|
"mkdocs-include-markdown-plugin",
|
|
38
46
|
"pymdown-extensions",
|
|
39
47
|
"pygments",
|
|
40
|
-
"pymdown-extensions",
|
|
41
48
|
"chex>=0.1.86",
|
|
42
49
|
"pre-commit",
|
|
43
50
|
]
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "1.4.dev367"
|
|
@@ -1 +0,0 @@
|
|
|
1
|
-
__version__ = "1.4.dev366"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|