haliax 1.4.dev303__tar.gz → 1.4.dev306__tar.gz

{haliax-1.4.dev303 → haliax-1.4.dev306}/.github/workflows/run_tests.yaml

@@ -17,7 +17,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install flake8 pytest
-          pip install --upgrade "jax[cpu]==0.4.25" "jaxlib[cpu]==0.4.25"
+          pip install --upgrade "jax[cpu]==0.4.30" "jaxlib[cpu]==0.4.30"
           pip install .[dev]
       - name: Test with pytest
         run: |

{haliax-1.4.dev303 → haliax-1.4.dev306}/PKG-INFO

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

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

{haliax-1.4.dev303 → haliax-1.4.dev306}/src/haliax/partitioning.py

@@ -128,22 +128,27 @@ def shard(x: T, mapping: Optional[ResourceMapping] = None, mesh: Optional[Mesh]
         warnings.warn("Sharding constraints are not supported in jit on metal", RuntimeWarning)
         return x
 
-    def _do_device_put(x):
-        if not isinstance(x, NamedArray):
-            return x
+    def _do_device_put(named):
+        if not isinstance(named, NamedArray):
+            return named
 
-        if not is_jax_array_like(x.array):
+        if not is_jax_array_like(named.array):
             # this happens when we filter out params for things like lora.
             # could use eqx.partition to avoid this, but eh
-            return x
+            return named
 
-        sharding = infer_resource_partitions(x, mapping, mesh=mesh, preserve_existing_shardings=False)
+        sharding = infer_resource_partitions(named, mapping, mesh=mesh, preserve_existing_shardings=False)
         assert isinstance(sharding, NamedSharding)
         if is_in_jit():
-            return with_sharding_constraint(x, sharding)
+            return with_sharding_constraint(named, sharding)
+        # as a special case, SingleDeviceShardings are routed through jit
+        elif isinstance(named.array.sharding, SingleDeviceSharding):
+            # TODO(dlwh): this should be unnecessary in JAX soon. Check after 2024-08-01
+            sharded_array = jax.jit(lambda x: x, out_shardings=sharding)(named)
+            return sharded_array
         else:
-            sharded_array = jax.device_put(x.array, sharding)
-            return NamedArray(sharded_array, x.axes)
+            ret = jax.device_put(named, sharding)
+            return ret
 
     return htu.tree_map(_do_device_put, x)
 
@@ -309,7 +314,7 @@ class _NamedJitWrapper(eqx.Module):
             output_shape = _cached_filter_eval_shape(self._fn, *args, **kwargs)
             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 not None:
                 in_resources = infer_resource_partitions(
                     (dynamic_donated, dynamic_reserved),
                     in_axis_resources,

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

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