haliax 1.4.dev391__tar.gz → 1.4.dev392__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.dev391 → haliax-1.4.dev392}/AGENTS.md +6 -8
- {haliax-1.4.dev391 → haliax-1.4.dev392}/PKG-INFO +1 -1
- haliax-1.4.dev392/src/haliax/__about__.py +1 -0
- haliax-1.4.dev391/src/haliax/__about__.py +0 -1
- {haliax-1.4.dev391 → haliax-1.4.dev392}/.coveragerc +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/.flake8 +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/.github/workflows/publish_dev.yaml +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/.github/workflows/run_pre_commit.yaml +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/.github/workflows/run_quick_levanter_tests.yaml +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/.github/workflows/run_tests.yaml +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/.gitignore +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/.playbooks/add-types.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/.playbooks/wrap-non-named.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/.pre-commit-config.yaml +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/.readthedocs.yaml +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/CONTRIBUTING.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/LICENSE +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/README.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/api.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/broadcasting.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/cheatsheet.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/css/material.css +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/css/mkdocstrings.css +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/faq.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/figures/data_parallel_mesh.png +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/figures/data_parallel_mesh_replicated.png +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/figures/device_mesh_1d.png +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/figures/device_mesh_1d_zero.png +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/figures/device_mesh_2d.png +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/figures/device_mesh_2d_batch_partitioned.png +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/figures/device_mesh_2d_data_replicated.png +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/figures/device_mesh_2d_data_replicated_mlp_partitioned.png +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/figures/device_mesh_2d_intermediate_fully_partitioned.png +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/figures/device_mesh_2d_zero.png +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/fp8.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/index.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/indexing.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/matmul.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/nn.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/partitioning.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/rearrange.ipynb +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/rearrange.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/requirements.txt +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/scan.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/state-dict.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/tutorial.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/typing.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/docs/vmap.md +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/mkdocs.yml +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/pyproject.toml +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/__init__.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/_src/__init__.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/_src/compile_utils.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/_src/dot.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/_src/einsum.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/_src/fp8.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/_src/parsing.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/_src/rearrange.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/_src/scan.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/_src/state_dict.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/_src/util.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/axis.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/core.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/debug.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/haxtyping.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/hof.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/jax_utils.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/nn/__init__.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/nn/activations.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/nn/attention.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/nn/conv.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/nn/dropout.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/nn/embedding.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/nn/linear.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/nn/loss.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/nn/mlp.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/nn/normalization.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/nn/pool.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/nn/scan.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/ops.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/partitioning.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/quantization.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/random.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/specialized_fns.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/state_dict.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/tree_util.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/types.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/util.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/src/haliax/wrap.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/core_test.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_attention.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_axis.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_conv.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_debug.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_dot.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_dtype_typing.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_einsum.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_fp8.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_hof.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_int8.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_namedarray_typing.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_nn.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_ops.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_parsing.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_partitioning.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_pool.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_random.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_rearrange.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_scan.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_scatter_gather.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_specialized_fns.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_state_dict.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_tree_util.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_utils.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/tests/test_visualize_sharding.py +0 -0
- {haliax-1.4.dev391 → haliax-1.4.dev392}/uv.lock +0 -0
|
@@ -17,7 +17,7 @@ repository. Follow these notes when implementing new features or fixing bugs.
|
|
|
17
17
|
## Playbook
|
|
18
18
|
|
|
19
19
|
- Adding Haliax-style tensor typing annotations are described in @.playbooks/add-types.md
|
|
20
|
-
- Wrapping standard JAX functions so they operate on `NamedArray`
|
|
20
|
+
- [Wrapping standard JAX functions](.playbooks/wrap-non-named.md) so they operate on `NamedArray`
|
|
21
21
|
|
|
22
22
|
## Code Style
|
|
23
23
|
|
|
@@ -58,15 +58,14 @@ repository. Follow these notes when implementing new features or fixing bugs.
|
|
|
58
58
|
|
|
59
59
|
* **Generic code**: many utilities are written with Python generics and dataclasses. Where possible,
|
|
60
60
|
write reusable functions or classes that operate over TypeVars instead of hard coding concrete types.
|
|
61
|
-
* **
|
|
62
|
-
declarative and typed.
|
|
63
|
-
* **Reproducibility**: Levanter aims for deterministic training where possible. Avoid sources of
|
|
61
|
+
* **Reproducibility**: Haliax aims for determinism where possible. Avoid sources of
|
|
64
62
|
nondeterminism unless explicitly required.
|
|
65
63
|
* Prefer Stacked with fold or scan over writing custom loops, for better compile times and gradient checkpointing support
|
|
64
|
+
* For configuration, we prefer frozen dataclasses over dictionaries.
|
|
66
65
|
|
|
67
66
|
## Library conventions
|
|
68
|
-
- Haliax revolves around `NamedArray` and
|
|
69
|
-
axes or axis names rather than hard‑coding positional dimensions.
|
|
67
|
+
- Haliax revolves around `NamedArray` and named shapes, either via Axis objects or "shape dicts" (e.g. `{"batch": 42, "embed": 16}).
|
|
68
|
+
Prefer APIs that accept axes or axis names rather than hard‑coding positional dimensions. In particular, use AxisSpec and AxisSelection where possible.
|
|
70
69
|
- Utilities should be written so they work with arbitrary axis names. Avoid relying on
|
|
71
70
|
fixed axis orders when possible.
|
|
72
71
|
- Use the provided modules in `haliax.nn` or Equinox when building neural network layers.
|
|
@@ -74,5 +73,4 @@ repository. Follow these notes when implementing new features or fixing bugs.
|
|
|
74
73
|
for a float32 array with a "batch" axis, or `ht.Float[NamedArray, "batch"]` for any floating point dtype.
|
|
75
74
|
|
|
76
75
|
## Documentation
|
|
77
|
-
- Public functions and modules require docstrings. If behavior is non‑obvious,
|
|
78
|
-
add examples in `docs/`.
|
|
76
|
+
- Public functions and modules require docstrings. If behavior is non‑obvious, add examples in `docs/`.
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: haliax
|
|
3
|
-
Version: 1.4.
|
|
3
|
+
Version: 1.4.dev392
|
|
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/
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "1.4.dev392"
|
|
@@ -1 +0,0 @@
|
|
|
1
|
-
__version__ = "1.4.dev391"
|
|
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
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|