haliax 1.4.dev300__tar.gz → 1.4.dev302__tar.gz

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

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

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