haliax 1.4.dev299__tar.gz → 1.4.dev300__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.dev299 → haliax-1.4.dev300}/.github/workflows/run_pre_commit.yaml +1 -1
- haliax-1.4.dev300/.github/workflows/run_quick_levanter_tests.yaml +36 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/.github/workflows/run_tests.yaml +3 -6
- {haliax-1.4.dev299 → haliax-1.4.dev300}/PKG-INFO +1 -1
- haliax-1.4.dev300/src/haliax/__about__.py +1 -0
- haliax-1.4.dev299/src/haliax/__about__.py +0 -1
- {haliax-1.4.dev299 → haliax-1.4.dev300}/.coveragerc +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/.flake8 +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/.github/workflows/publish_dev.yaml +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/.gitignore +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/.pre-commit-config.yaml +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/.readthedocs.yaml +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/CONTRIBUTING.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/LICENSE +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/README.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/api.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/broadcasting.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/cheatsheet.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/css/material.css +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/css/mkdocstrings.css +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/faq.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/data_parallel_mesh.png +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/data_parallel_mesh_replicated.png +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_1d.png +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_1d_zero.png +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_2d.png +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_2d_batch_partitioned.png +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_2d_data_replicated.png +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_2d_data_replicated_mlp_partitioned.png +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_2d_intermediate_fully_partitioned.png +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_2d_zero.png +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/fp8.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/hof.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/index.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/indexing.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/matmul.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/nn.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/partitioning.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/rearrange.ipynb +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/rearrange.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/requirements.txt +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/tutorial.md +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/mkdocs.yml +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/pyproject.toml +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/__init__.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/__init__.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/compile_utils.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/dot.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/einsum.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/fp8.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/parsing.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/rearrange.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/util.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/axis.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/core.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/debug.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/hof.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/jax_utils.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/__init__.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/activations.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/attention.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/conv.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/dropout.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/embedding.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/linear.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/loss.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/mlp.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/normalization.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/pool.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/scan.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/ops.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/partitioning.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/quantization.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/random.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/specialized_fns.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/tree_util.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/types.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/util.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/wrap.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/core_test.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_attention.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_axis.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_conv.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_debug.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_dot.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_einsum.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_fp8.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_hof.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_nn.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_ops.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_parsing.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_partitioning.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_pool.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_random.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_rearrange.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_scan.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_specialized_fns.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_tree_util.py +0 -0
- {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_utils.py +0 -0
|
@@ -0,0 +1,36 @@
|
|
|
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 --upgrade "jax[cpu]==0.4.26" "jaxlib[cpu]==0.4.26"
|
|
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 .
|
|
28
|
+
- name: Install Haliax on top
|
|
29
|
+
run: |
|
|
30
|
+
# install second since levanter will install a built version of haliax
|
|
31
|
+
cd ../haliax
|
|
32
|
+
pip install .[dev]
|
|
33
|
+
- name: Test levanter with pytest
|
|
34
|
+
run: |
|
|
35
|
+
cd ../levanter
|
|
36
|
+
XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:../src pytest tests -m "not entry and not slow"
|
|
@@ -1,21 +1,18 @@
|
|
|
1
1
|
name: Run Tests
|
|
2
2
|
|
|
3
|
-
on: [push]
|
|
3
|
+
on: [push, pull_request]
|
|
4
4
|
|
|
5
5
|
jobs:
|
|
6
6
|
build:
|
|
7
7
|
|
|
8
8
|
runs-on: ubuntu-latest
|
|
9
|
-
strategy:
|
|
10
|
-
matrix:
|
|
11
|
-
python-version: ["3.10.11"]
|
|
12
9
|
|
|
13
10
|
steps:
|
|
14
11
|
- uses: actions/checkout@v3
|
|
15
|
-
- name: Set up Python
|
|
12
|
+
- name: Set up Python 3.10.11
|
|
16
13
|
uses: actions/setup-python@v4
|
|
17
14
|
with:
|
|
18
|
-
python-version:
|
|
15
|
+
python-version: 3.10.11
|
|
19
16
|
- name: Install dependencies
|
|
20
17
|
run: |
|
|
21
18
|
python -m pip install --upgrade pip
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.3
|
|
2
2
|
Name: haliax
|
|
3
|
-
Version: 1.4.
|
|
3
|
+
Version: 1.4.dev300
|
|
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.dev300"
|
|
@@ -1 +0,0 @@
|
|
|
1
|
-
__version__ = "1.4.dev299"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|