haliax 1.4.dev374__tar.gz → 1.4.dev375__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.dev375/.github/workflows/run_quick_levanter_tests.yaml +48 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/PKG-INFO +1 -1
- haliax-1.4.dev375/src/haliax/__about__.py +1 -0
- haliax-1.4.dev374/.github/workflows/run_quick_levanter_tests.yaml +0 -38
- haliax-1.4.dev374/src/haliax/__about__.py +0 -1
- {haliax-1.4.dev374 → haliax-1.4.dev375}/.coveragerc +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/.flake8 +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/.github/workflows/publish_dev.yaml +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/.github/workflows/run_pre_commit.yaml +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/.github/workflows/run_tests.yaml +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/.gitignore +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/.playbooks/add-types.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/.pre-commit-config.yaml +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/.readthedocs.yaml +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/AGENTS.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/CONTRIBUTING.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/LICENSE +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/README.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/api.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/broadcasting.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/cheatsheet.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/css/material.css +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/css/mkdocstrings.css +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/faq.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/figures/data_parallel_mesh.png +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/figures/data_parallel_mesh_replicated.png +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/figures/device_mesh_1d.png +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/figures/device_mesh_1d_zero.png +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/figures/device_mesh_2d.png +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/figures/device_mesh_2d_batch_partitioned.png +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/figures/device_mesh_2d_data_replicated.png +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/figures/device_mesh_2d_data_replicated_mlp_partitioned.png +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/figures/device_mesh_2d_intermediate_fully_partitioned.png +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/figures/device_mesh_2d_zero.png +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/fp8.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/index.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/indexing.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/matmul.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/nn.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/partitioning.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/rearrange.ipynb +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/rearrange.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/requirements.txt +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/scan.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/state-dict.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/tutorial.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/typing.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/docs/vmap.md +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/mkdocs.yml +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/pyproject.toml +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/__init__.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/_src/__init__.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/_src/compile_utils.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/_src/dot.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/_src/einsum.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/_src/fp8.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/_src/parsing.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/_src/rearrange.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/_src/scan.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/_src/state_dict.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/_src/util.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/axis.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/core.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/debug.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/haxtyping.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/hof.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/jax_utils.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/nn/__init__.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/nn/activations.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/nn/attention.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/nn/conv.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/nn/dropout.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/nn/embedding.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/nn/linear.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/nn/loss.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/nn/mlp.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/nn/normalization.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/nn/pool.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/nn/scan.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/ops.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/partitioning.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/quantization.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/random.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/specialized_fns.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/state_dict.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/tree_util.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/types.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/util.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/src/haliax/wrap.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/core_test.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_attention.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_axis.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_conv.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_debug.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_dot.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_dtype_typing.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_einsum.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_fp8.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_hof.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_int8.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_namedarray_typing.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_nn.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_ops.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_parsing.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_partitioning.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_pool.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_random.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_rearrange.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_scan.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_scatter_gather.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_specialized_fns.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_state_dict.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_tree_util.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/tests/test_utils.py +0 -0
- {haliax-1.4.dev374 → haliax-1.4.dev375}/uv.lock +0 -0
|
@@ -0,0 +1,48 @@
|
|
|
1
|
+
name: Run Levanter & Haliax Tests
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
pull_request:
|
|
5
|
+
branches: [main]
|
|
6
|
+
|
|
7
|
+
jobs:
|
|
8
|
+
test:
|
|
9
|
+
runs-on: ubuntu-latest
|
|
10
|
+
strategy:
|
|
11
|
+
matrix:
|
|
12
|
+
python-version: [3.10, 3.11]
|
|
13
|
+
steps:
|
|
14
|
+
- name: Checkout repository
|
|
15
|
+
uses: actions/checkout@v3
|
|
16
|
+
|
|
17
|
+
- name: Cache pip and uv artifacts
|
|
18
|
+
uses: actions/cache@v3
|
|
19
|
+
with:
|
|
20
|
+
path: |
|
|
21
|
+
~/.cache/pip
|
|
22
|
+
~/.cache/uv
|
|
23
|
+
key: ${{ runner.os }}-pip-uv-${{ matrix.python-version }}-${{ hashFiles('**/requirements.lock', '**/uv.lock') }}
|
|
24
|
+
restore-keys: |
|
|
25
|
+
${{ runner.os }}-pip-uv-${{ matrix.python-version }}-
|
|
26
|
+
|
|
27
|
+
- name: Set up Python ${{ matrix.python-version }}
|
|
28
|
+
uses: actions/setup-python@v4
|
|
29
|
+
with:
|
|
30
|
+
python-version: ${{ matrix.python-version }}
|
|
31
|
+
|
|
32
|
+
- name: Install uv and sync dependencies
|
|
33
|
+
run: |
|
|
34
|
+
pip install uv
|
|
35
|
+
uv venv .venv
|
|
36
|
+
uv sync
|
|
37
|
+
|
|
38
|
+
- name: Install Levanter and Haliax in editable mode
|
|
39
|
+
run: |
|
|
40
|
+
uv pip install -e .[tests]
|
|
41
|
+
uv pip install -e ../haliax
|
|
42
|
+
|
|
43
|
+
- name: Run pytest suite
|
|
44
|
+
env:
|
|
45
|
+
XLA_FLAGS: --xla_force_host_platform_device_count=8
|
|
46
|
+
VIRTUAL_ENV: ${{ github.workspace }}/.venv
|
|
47
|
+
run: |
|
|
48
|
+
pytest tests -m "not entry and not slow"
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: haliax
|
|
3
|
-
Version: 1.4.
|
|
3
|
+
Version: 1.4.dev375
|
|
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.dev375"
|
|
@@ -1,38 +0,0 @@
|
|
|
1
|
-
name: Run Levanter Tests
|
|
2
|
-
|
|
3
|
-
on: [pull_request]
|
|
4
|
-
|
|
5
|
-
jobs:
|
|
6
|
-
build:
|
|
7
|
-
|
|
8
|
-
runs-on: ubuntu-latest
|
|
9
|
-
|
|
10
|
-
steps:
|
|
11
|
-
- uses: actions/checkout@v3
|
|
12
|
-
- name: Set up Python 3.10.11
|
|
13
|
-
uses: actions/setup-python@v4
|
|
14
|
-
with:
|
|
15
|
-
python-version: 3.10.11
|
|
16
|
-
- name: Install dependencies
|
|
17
|
-
run: |
|
|
18
|
-
python -m pip install --upgrade pip
|
|
19
|
-
pip install flake8 pytest
|
|
20
|
-
pip install "jax[cpu]==0.5.3" "jaxlib[cpu]==0.5.3" .[dev]
|
|
21
|
-
|
|
22
|
-
- name: Install Levanter from source
|
|
23
|
-
run: |
|
|
24
|
-
cd ..
|
|
25
|
-
git clone https://github.com/stanford-crfm/levanter.git
|
|
26
|
-
cd levanter
|
|
27
|
-
pip install -e .[tests]
|
|
28
|
-
pip install -r tests/requirements.txt
|
|
29
|
-
# i don't know why this is necessary
|
|
30
|
-
pip install tensorboardX
|
|
31
|
-
- name: Install Haliax on top
|
|
32
|
-
run: |
|
|
33
|
-
# install second since levanter will install a built version of haliax
|
|
34
|
-
cd ../haliax
|
|
35
|
-
- name: Test levanter with pytest
|
|
36
|
-
run: |
|
|
37
|
-
cd ../levanter
|
|
38
|
-
XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:../src pytest tests -m "not entry and not slow"
|
|
@@ -1 +0,0 @@
|
|
|
1
|
-
__version__ = "1.4.dev374"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|