haliax 1.4.dev373__tar.gz → 1.4.dev375__tar.gz

haliax-1.4.dev375/.github/workflows/run_quick_levanter_tests.yaml

@@ -0,0 +1,48 @@
+name: Run Levanter & Haliax Tests
+
+on:
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.10, 3.11]
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v3
+
+      - name: Cache pip and uv artifacts
+        uses: actions/cache@v3
+        with:
+          path: |
+            ~/.cache/pip
+            ~/.cache/uv
+          key: ${{ runner.os }}-pip-uv-${{ matrix.python-version }}-${{ hashFiles('**/requirements.lock', '**/uv.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-pip-uv-${{ matrix.python-version }}-
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install uv and sync dependencies
+        run: |
+          pip install uv
+          uv venv .venv
+          uv sync
+
+      - name: Install Levanter and Haliax in editable mode
+        run: |
+          uv pip install -e .[tests]
+          uv pip install -e ../haliax
+
+      - name: Run pytest suite
+        env:
+          XLA_FLAGS: --xla_force_host_platform_device_count=8
+          VIRTUAL_ENV: ${{ github.workspace }}/.venv
+        run: |
+          pytest tests -m "not entry and not slow"

{haliax-1.4.dev373 → haliax-1.4.dev375}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev373
+Version: 1.4.dev375
 Summary: Named Tensors for Legible Deep Learning in JAX
 Project-URL: Homepage, https://github.com/stanford-crfm/haliax
 Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/

haliax-1.4.dev375/src/haliax/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.4.dev375"

{haliax-1.4.dev373 → haliax-1.4.dev375}/src/haliax/core.py

@@ -27,6 +27,7 @@ from .axis import (
     PartialShapeDict,
     ShapeDict,
     axis_name,
+    axis_spec_to_shape_dict,
     axis_spec_to_tuple,
     dslice,
     eliminate_axes,
@@ -1679,7 +1680,7 @@ def broadcast_to(
     all_axes = to_add + a.axes
 
     if enforce_no_extra_axes and len(all_axes) != len(axes):
-        raise ValueError(f"Cannot broadcast {a} to {axes}: extra axes present")
+        raise ValueError(f"Cannot broadcast {a.shape} to {axis_spec_to_shape_dict(axes)}: extra axes present")
 
     extra_axes = tuple(ax for ax in a.axes if ax not in axes)
 

haliax-1.4.dev373/.github/workflows/run_quick_levanter_tests.yaml

@@ -1,38 +0,0 @@
-name: Run Levanter Tests
-
-on: [pull_request]
-
-jobs:
-  build:
-
-    runs-on: ubuntu-latest
-
-    steps:
-      - uses: actions/checkout@v3
-      - name: Set up Python 3.10.11
-        uses: actions/setup-python@v4
-        with:
-          python-version: 3.10.11
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install flake8 pytest
-          pip install "jax[cpu]==0.5.3" "jaxlib[cpu]==0.5.3" .[dev]
-
-      - name: Install Levanter from source
-        run: |
-          cd ..
-          git clone https://github.com/stanford-crfm/levanter.git
-          cd levanter
-          pip install -e .[tests]
-          pip install -r tests/requirements.txt
-          # i don't know why this is necessary
-          pip install tensorboardX
-      - name: Install Haliax on top
-        run: |
-          # install second since levanter will install a built version of haliax
-          cd ../haliax
-      - name: Test levanter with pytest
-        run: |
-          cd ../levanter
-          XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:../src pytest tests -m "not entry and not slow"

haliax-1.4.dev373/src/haliax/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.4.dev373"