haliax 1.4.dev374__tar.gz → 1.4.dev378__tar.gz

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

@@ -0,0 +1,49 @@
+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
+          uv pip install pytest
+
+      - name: Run pytest suite
+        env:
+          XLA_FLAGS: --xla_force_host_platform_device_count=8
+          VIRTUAL_ENV: ${{ github.workspace }}/.venv
+        run: |
+          uv run pytest tests -m "not entry and not slow"

{haliax-1.4.dev374 → haliax-1.4.dev378}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev374
+Version: 1.4.dev378
 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.dev378/src/haliax/__about__.py

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

{haliax-1.4.dev374 → haliax-1.4.dev378}/src/haliax/nn/dropout.py

@@ -61,8 +61,8 @@ class Dropout(eqx.Module):
     """
 
     # key difference from equinox: these are static fields
-    pdrop: float = eqx.static_field()
-    broadcast_axes: Optional[AxisSpec] = eqx.static_field()
+    pdrop: float = eqx.field(static=True)
+    broadcast_axes: Optional[AxisSpec] = eqx.field(static=True)
     inference: bool = False  # note: not static
 
     def __init__(

{haliax-1.4.dev374 → haliax-1.4.dev378}/src/haliax/nn/embedding.py

@@ -17,8 +17,8 @@ class Embedding(eqx.Module):
     weight: NamedArray
 
     # axes
-    Vocab: Axis = eqx.static_field()
-    Embed: AxisSpec = eqx.static_field()
+    Vocab: Axis = eqx.field(static=True)
+    Embed: AxisSpec = eqx.field(static=True)
 
     @staticmethod
     def init(Vocab: Axis, Embed: AxisSpec, *, init_scale: float = 1, key, initializer_range: Optional[float] = None):

{haliax-1.4.dev374 → haliax-1.4.dev378}/src/haliax/nn/linear.py

@@ -28,8 +28,8 @@ class Linear(ModuleWithStateDictSerialization):
     weight: NamedArray
     bias: Optional[NamedArray]
 
-    In: AxisSpec = eqx.static_field()
-    Out: AxisSpec = eqx.static_field()
+    In: AxisSpec = eqx.field(static=True)
+    Out: AxisSpec = eqx.field(static=True)
     dot_general: DotGeneralOp = eqx.field(default_factory=DotGeneralOp.default)
 
     @staticmethod
@@ -140,9 +140,9 @@ class MoELinear(eqx.Module):
     weight: NamedArray
     bias: Optional[NamedArray]
 
-    Experts: AxisSpec = eqx.static_field()
-    In: Axis = eqx.static_field()
-    Out: Axis = eqx.static_field()
+    Experts: AxisSpec = eqx.field(static=True)
+    In: Axis = eqx.field(static=True)
+    Out: Axis = eqx.field(static=True)
     # TODO: support quanitization for ragged_dot?
     # dot_general: DotGeneralOp = eqx.field(default_factory=DotGeneralOp.default)
 

{haliax-1.4.dev374 → haliax-1.4.dev378}/src/haliax/nn/normalization.py

@@ -20,10 +20,10 @@ A = TypeVar("A", Scalar, NamedArray, jnp.ndarray)
 
 
 class LayerNormBase(ModuleWithStateDictSerialization):
-    axis: AxisSpec = eqx.static_field()
+    axis: AxisSpec = eqx.field(static=True)
     weight: Optional[NamedArray]
     bias: Optional[NamedArray]
-    eps: float = eqx.static_field(default=1e-5)
+    eps: float = eqx.field(default=1e-5, static=True)
     dtype: Optional[jnp.dtype] = eqx.field(default=None, static=True)
 
     @abstractmethod

{haliax-1.4.dev374 → haliax-1.4.dev378}/src/haliax/nn/scan.py

@@ -90,8 +90,8 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
     """
 
     blocks: Sequence[M]
-    Block: Axis = eqx.static_field()
-    gradient_checkpointing: ScanCheckpointPolicy = eqx.static_field()
+    Block: Axis = eqx.field(static=True)
+    gradient_checkpointing: ScanCheckpointPolicy = eqx.field(static=True)
 
     @classmethod
     def init(
@@ -279,8 +279,8 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
     # TODO: we can probably make this module support pipeline parallelism, but that's a whole project in itself
 
     stacked: M
-    Block: Axis = eqx.static_field()
-    gradient_checkpointing: ScanCheckpointPolicy = eqx.static_field()
+    Block: Axis = eqx.field(static=True)
+    gradient_checkpointing: ScanCheckpointPolicy = eqx.field(static=True)
 
     @classmethod
     def init(

haliax-1.4.dev374/.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.dev374/src/haliax/__about__.py

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