haliax 1.4.dev402__tar.gz → 1.4.dev403__tar.gz

{haliax-1.4.dev402 → haliax-1.4.dev403}/PKG-INFO

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

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

{haliax-1.4.dev402 → haliax-1.4.dev403}/src/haliax/partitioning.py

@@ -1,4 +1,5 @@
 import contextlib
+import dataclasses
 import functools
 import threading
 import typing
@@ -210,6 +211,37 @@ def pspec_for(
                 return None
             else:
                 return pspec_for_axis(node.axes, resource_mapping)
+        elif isinstance(node, eqx.Module):
+            # handle eqx.Module explicitly so that we can look at axis_names metadata
+            updates: dict[str, typing.Any] = {}
+            for field in dataclasses.fields(node):
+                if field.metadata.get("static", False):
+                    continue
+
+                value = getattr(node, field.name)
+                axis_names = field.metadata.get("axis_names") if field.metadata is not None else None
+                if axis_names is not None and is_jax_array_like(value):
+                    current_sharding = (
+                        getattr(value, "sharding", None) if preserve_existing_shardings else None
+                    )
+                    if current_sharding is not None:
+                        updates[field.name] = None
+                    else:
+                        updates[field.name] = pspec_for_axis(axis_names, resource_mapping)
+                else:
+                    updates[field.name] = htu.tree_map(
+                        partition_spec, value, is_leaf=lambda x: isinstance(x, eqx.Module)
+                    )
+
+            new_node = object.__new__(type(node))
+            for field in dataclasses.fields(node):
+                object.__setattr__(
+                    new_node,
+                    field.name,
+                    updates.get(field.name, getattr(node, field.name)),
+                )
+
+            return new_node
         elif is_jax_array_like(node):
             sharding = getattr(node, "sharding", None)
             # TODO: these are usually replicated. Is there a better way to tell?
@@ -232,7 +264,7 @@ def pspec_for(
         else:
             return None
 
-    return htu.tree_map(partition_spec, tree)
+    return htu.tree_map(partition_spec, tree, is_leaf=lambda x: isinstance(x, eqx.Module))
 
 
 def infer_resource_partitions(

{haliax-1.4.dev402 → haliax-1.4.dev403}/tests/test_partitioning.py

@@ -59,6 +59,34 @@ def test_pspec_for_named_axes():
         assert specs.unnamed1 == PartitionSpec(None)
 
 
+class ArrayModule(eqx.Module):
+    arr: Array = hax.field(axis_names=("dim2", "dim3"))
+
+
+def test_pspec_for_plain_array_axis_names():
+    mesh = Mesh(np.array(jax.devices()).reshape(-1, 1), (ResourceAxis.DATA, ResourceAxis.MODEL))
+    with axis_mapping(resource_map), mesh:
+        mod = ArrayModule(jnp.ones((Dim2.size, Dim3.size)))
+
+        specs: ArrayModule = pspec_for(mod, preserve_existing_shardings=False)
+
+        assert specs.arr == PartitionSpec(ResourceAxis.DATA, ResourceAxis.MODEL)
+
+
+class NestedArrayModule(eqx.Module):
+    inner: ArrayModule
+
+
+def test_pspec_for_plain_array_axis_names_nested_module():
+    mesh = Mesh(np.array(jax.devices()).reshape(-1, 1), (ResourceAxis.DATA, ResourceAxis.MODEL))
+    with axis_mapping(resource_map), mesh:
+        mod = NestedArrayModule(ArrayModule(jnp.ones((Dim2.size, Dim3.size))))
+
+        specs: NestedArrayModule = pspec_for(mod, preserve_existing_shardings=False)
+
+        assert specs.inner.arr == PartitionSpec(ResourceAxis.DATA, ResourceAxis.MODEL)
+
+
 class MyModuleInit(eqx.Module):
     named: NamedArray
     unnamed1: Array

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

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