haliax 1.4.dev375__tar.gz → 1.4.dev378__tar.gz

{haliax-1.4.dev375 → haliax-1.4.dev378}/.github/workflows/run_quick_levanter_tests.yaml

@@ -9,7 +9,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.10, 3.11]
+        python-version: ["3.10", "3.11"]
     steps:
       - name: Checkout repository
         uses: actions/checkout@v3
@@ -39,10 +39,11 @@ jobs:
         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: |
-          pytest tests -m "not entry and not slow"
+          uv run pytest tests -m "not entry and not slow"

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

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

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