haliax 1.4.dev298__tar.gz → 1.4.dev302__tar.gz

{haliax-1.4.dev298 → haliax-1.4.dev302}/.github/workflows/run_pre_commit.yaml

@@ -1,6 +1,6 @@
 name: Pre-Commit
 
-on: [push]
+on: [push, pull_request]
 
 jobs:
   build:

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

@@ -0,0 +1,36 @@
+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 --upgrade "jax[cpu]==0.4.26" "jaxlib[cpu]==0.4.26"
+
+      - name: Install Levanter from source
+        run: |
+          cd ..
+          git clone https://github.com/stanford-crfm/levanter.git
+          cd levanter
+          pip install -e .
+      - name: Install Haliax on top
+        run: |
+          # install second since levanter will install a built version of haliax
+          cd ../haliax
+          pip install .[dev]
+      - 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.dev298 → haliax-1.4.dev302}/.github/workflows/run_tests.yaml

@@ -1,21 +1,18 @@
 name: Run Tests
 
-on: [push]
+on: [push, pull_request]
 
 jobs:
   build:
 
     runs-on: ubuntu-latest
-    strategy:
-      matrix:
-        python-version: ["3.10.11"]
 
     steps:
       - uses: actions/checkout@v3
-      - name: Set up Python ${{ matrix.python-version }}
+      - name: Set up Python 3.10.11
         uses: actions/setup-python@v4
         with:
-          python-version: ${{ matrix.python-version }}
+          python-version: 3.10.11
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip

{haliax-1.4.dev298 → haliax-1.4.dev302}/PKG-INFO

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

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

{haliax-1.4.dev298 → haliax-1.4.dev302}/src/haliax/partitioning.py

@@ -141,9 +141,13 @@ def shard(x: T, mapping: Optional[ResourceMapping] = None, mesh: Optional[Mesh]
         assert isinstance(sharding, NamedSharding)
         if is_in_jit():
             return with_sharding_constraint(x, sharding)
-        else:
+        elif sharding.is_fully_addressable:
             sharded_array = jax.device_put(x.array, sharding)
             return NamedArray(sharded_array, x.axes)
+        else:
+            # sharded_array = jax.device_put(x.array, sharding)
+            ret = eqx.filter_jit(lambda x: with_sharding_constraint(x, sharding))(x)
+            return ret
 
     return htu.tree_map(_do_device_put, x)
 
@@ -310,8 +314,6 @@ class _NamedJitWrapper(eqx.Module):
             my_pjit_args = dict(**self._pjit_args)
 
             if in_axis_resources is not None or axis_resources is not None:
-                if in_axis_resources is None:
-                    in_axis_resources = axis_resources
                 in_resources = infer_resource_partitions(
                     (dynamic_donated, dynamic_reserved),
                     in_axis_resources,

{haliax-1.4.dev298 → haliax-1.4.dev302}/tests/test_partitioning.py

@@ -99,7 +99,7 @@ def test_pjit_class_init_with_args():
 
         devices = jax.devices()
         with Mesh(np.array(devices).reshape(-1, 1), (ResourceAxis.DATA, ResourceAxis.MODEL)):
-            mod = named_jit(ModWithArgs)(hax.ones((Dim1, Dim2)))
+            mod = named_jit(ModWithArgs)(hax.shard(hax.ones((Dim1, Dim2))))
         assert isinstance(mod, ModWithArgs)
         assert mod.array.array.shape == (Dim1.size, Dim2.size)
         assert mod.array2.array.shape == (Dim3.size,)
@@ -173,7 +173,7 @@ def test_shard_with_axis_mapping_inside_jit():
 
             jax.debug.inspect_array_sharding(arr.array, callback=lambda x: assert_eq(x, expected))
 
-        @named_jit(in_axis_resources={}, out_axis_resources=resource_map)
+        @named_jit(out_axis_resources=resource_map)
         def do_shard(x, y):
             x = hax.shard(x, resource_map)
             assert_inside_pjit(x, NamedSharding(mesh, PartitionSpec(None, ResourceAxis.DATA)))
@@ -293,3 +293,25 @@ def test_cross_device_sharding():
         z_devices = z.array.devices()
 
         assert set(d.platform for d in x_devices) == set(d.platform for d in z_devices)
+
+
+def test_named_jit_no_in_axis_resources():
+    mesh = Mesh(np.array(jax.devices()).reshape(-1, 1), (ResourceAxis.DATA, ResourceAxis.MODEL))
+    with axis_mapping(resource_map), mesh:
+
+        class MyModule(eqx.Module):
+            array: NamedArray
+
+            def __init__(self):
+                self.array = hax.ones((Dim1, Dim2))
+
+        data = hax.ones((Dim1, Dim2))
+        data = hax.shard(data, {})
+
+        @named_jit(axis_resources=resource_map)
+        def fn(data):
+            mod = MyModule()
+            return mod.array
+
+        r = fn(data)
+        assert r.array.sharding.device_set == set(jax.devices())

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

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

{haliax-1.4.dev298 → haliax-1.4.dev302}/src/haliax/__init__.py

@@ -56,8 +56,8 @@ from .core import (
     unflatten_axis,
     updated_slice,
 )
-from .jax_utils import filter_checkpoint
 from .hof import fold, map, scan, vmap
+from .jax_utils import filter_checkpoint
 from .ops import clip, isclose, pad_left, trace, tril, triu, where
 from .partitioning import auto_sharded, axis_mapping, fsdp, named_jit, shard, shard_with_axis_mapping
 from .specialized_fns import top_k