haliax 1.4.dev367__tar.gz → 1.4.dev368__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.dev368/.playbooks/add-types.md +45 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/PKG-INFO +1 -1
- haliax-1.4.dev368/src/haliax/__about__.py +1 -0
- haliax-1.4.dev367/src/haliax/__about__.py +0 -1
- {haliax-1.4.dev367 → haliax-1.4.dev368}/.coveragerc +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/.flake8 +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/.github/workflows/publish_dev.yaml +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/.github/workflows/run_pre_commit.yaml +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/.github/workflows/run_quick_levanter_tests.yaml +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/.github/workflows/run_tests.yaml +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/.gitignore +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/.pre-commit-config.yaml +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/.readthedocs.yaml +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/AGENTS.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/CONTRIBUTING.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/LICENSE +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/README.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/api.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/broadcasting.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/cheatsheet.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/css/material.css +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/css/mkdocstrings.css +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/faq.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/figures/data_parallel_mesh.png +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/figures/data_parallel_mesh_replicated.png +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/figures/device_mesh_1d.png +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/figures/device_mesh_1d_zero.png +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/figures/device_mesh_2d.png +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/figures/device_mesh_2d_batch_partitioned.png +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/figures/device_mesh_2d_data_replicated.png +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/figures/device_mesh_2d_data_replicated_mlp_partitioned.png +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/figures/device_mesh_2d_intermediate_fully_partitioned.png +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/figures/device_mesh_2d_zero.png +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/fp8.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/index.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/indexing.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/matmul.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/nn.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/partitioning.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/rearrange.ipynb +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/rearrange.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/requirements.txt +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/scan.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/state-dict.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/tutorial.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/typing.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/docs/vmap.md +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/mkdocs.yml +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/pyproject.toml +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/__init__.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/_src/__init__.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/_src/compile_utils.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/_src/dot.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/_src/einsum.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/_src/fp8.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/_src/parsing.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/_src/rearrange.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/_src/scan.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/_src/state_dict.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/_src/util.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/axis.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/core.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/debug.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/haxtyping.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/hof.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/jax_utils.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/nn/__init__.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/nn/activations.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/nn/attention.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/nn/conv.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/nn/dropout.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/nn/embedding.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/nn/linear.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/nn/loss.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/nn/mlp.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/nn/normalization.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/nn/pool.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/nn/scan.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/ops.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/partitioning.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/quantization.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/random.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/specialized_fns.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/state_dict.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/tree_util.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/types.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/util.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/src/haliax/wrap.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/core_test.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_attention.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_axis.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_conv.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_debug.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_dot.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_dtype_typing.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_einsum.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_fp8.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_hof.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_int8.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_namedarray_typing.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_nn.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_ops.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_parsing.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_partitioning.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_pool.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_random.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_rearrange.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_scan.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_scatter_gather.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_specialized_fns.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_state_dict.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_tree_util.py +0 -0
- {haliax-1.4.dev367 → haliax-1.4.dev368}/tests/test_utils.py +0 -0
|
@@ -0,0 +1,45 @@
|
|
|
1
|
+
# Adding NamedArray Type Annotations
|
|
2
|
+
|
|
3
|
+
This playbook explains how to migrate existing code to use the new type
|
|
4
|
+
annotation helpers described in `docs/typing.md`.
|
|
5
|
+
|
|
6
|
+
1. **Read the documentation**: Familiarise yourself with `docs/typing.md`.
|
|
7
|
+
It describes the `Named[...]` syntax and the dtype aware helpers in
|
|
8
|
+
`haliax.typing`.
|
|
9
|
+
2. **Annotate parameters**: Replace plain `NamedArray` annotations with
|
|
10
|
+
`Named[...]` that lists the required axes. Use ellipses or sets when the
|
|
11
|
+
exact order is flexible.
|
|
12
|
+
|
|
13
|
+
```python
|
|
14
|
+
from haliax import Named
|
|
15
|
+
|
|
16
|
+
# old
|
|
17
|
+
def foo(x: NamedArray) -> NamedArray:
|
|
18
|
+
...
|
|
19
|
+
|
|
20
|
+
# new
|
|
21
|
+
def foo(x: Named["batch", "embed"]) -> Named["batch", "embed"]:
|
|
22
|
+
...
|
|
23
|
+
```
|
|
24
|
+
3. **Annotate dtypes when needed**: If the dtype matters, import symbolic dtypes
|
|
25
|
+
from `haliax.typing` (e.g. `ht.f32`, `ht.i32`). They can be indexed just
|
|
26
|
+
like `Named`.
|
|
27
|
+
|
|
28
|
+
```python
|
|
29
|
+
import haliax.typing as ht
|
|
30
|
+
|
|
31
|
+
def bar(x: ht.f32["batch"]):
|
|
32
|
+
...
|
|
33
|
+
```
|
|
34
|
+
4. **Runtime validation**: Use `arr.matches_axes(...)` to check that a
|
|
35
|
+
`NamedArray` conforms to the expected axes and dtype at runtime.
|
|
36
|
+
|
|
37
|
+
```python
|
|
38
|
+
if not arr.matches_axes(Named["batch embed ..."]):
|
|
39
|
+
raise ValueError("unexpected axes")
|
|
40
|
+
```
|
|
41
|
+
5. **Update return types**: Functions returning `NamedArray` should annotate
|
|
42
|
+
their return values using the same conventions.
|
|
43
|
+
|
|
44
|
+
Following these steps will gradually port legacy code to the new
|
|
45
|
+
annotation style.
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: haliax
|
|
3
|
-
Version: 1.4.
|
|
3
|
+
Version: 1.4.dev368
|
|
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.dev368"
|
|
@@ -1 +0,0 @@
|
|
|
1
|
-
__version__ = "1.4.dev367"
|
|
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
|