haliax 1.4.dev400__tar.gz → 1.4.dev403__tar.gz

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

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

@@ -17,6 +17,7 @@ import haliax.random as random
 import haliax.state_dict as state_dict
 import haliax.tree_util as tree_util
 import haliax.util as util
+from .field import field
 
 from ._src.dot import dot
 from ._src.einsum import einsum
@@ -932,6 +933,7 @@ __all__ = [
     "tree_util",
     "nn",
     "state_dict",
+    "field",
     "Axis",
     "AxisSpec",
     "AxisSelection",

haliax-1.4.dev403/src/haliax/field.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from typing import Any, Callable
+
+import equinox as eqx
+
+
+def field(
+    *,
+    converter: Callable[[Any], Any] | None = None,
+    static: bool = False,
+    axis_names: tuple[str, ...] | None = None,
+    **kwargs,
+):
+    """Wrapper around :func:`equinox.field` with optional ``axis_names`` metadata.
+
+    Args:
+        converter: Optional function applied to the value during dataclass initialisation.
+        static: Whether the field is static in the PyTree.
+        axis_names: Optional axis names associated with array fields. Cannot be
+            specified together with ``static=True``.
+        **kwargs: Additional keyword arguments forwarded to :func:`dataclasses.field`.
+
+    Returns:
+        A dataclasses field configured like :func:`equinox.field` with additional
+        ``axis_names`` metadata.
+    """
+    if static and axis_names is not None:
+        raise ValueError("axis_names cannot be specified together with static=True")
+
+    metadata = dict(kwargs.pop("metadata", {}))
+    metadata["axis_names"] = axis_names
+
+    field_kwargs = {}
+    if converter is not None:
+        field_kwargs["converter"] = converter
+
+    return eqx.field(static=static, metadata=metadata, **field_kwargs, **kwargs)

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

@@ -1,4 +1,5 @@
 import contextlib
+import dataclasses
 import functools
 import threading
 import typing
@@ -177,24 +178,21 @@ def shard_with_axis_mapping(x: T, mapping: ResourceMapping, mesh: Optional[Mesh]
     return shard(x, mapping, mesh)
 
 
-def infer_resource_partitions(
+def pspec_for(
     tree: PyTree,
     resource_mapping: Optional[ResourceMapping] = None,
     preserve_existing_shardings: bool = True,
     use_auto_sharding: bool = True,
-    mesh: Optional[Mesh] = None,
 ) -> PyTree:
-    """
-    Infer the sharding for a module, to be used with named_jit.
-    The basic idea is to tree all NamedArrays as leaves for the purposes of this function,
-    and to create NamedShardings from those names plus the resource_mapping.
-    If preserve_existing_shardings is True, then NamedArrays that are already sharded are left alone.
+    """Infer the :class:`PartitionSpec` for a module.
 
-    If resource_mapping is not provided, this function attempts to use the global resource mapping.
+    This behaves like :func:`infer_resource_partitions` but returns ``PartitionSpec``
+    objects instead of :class:`~jax.sharding.NamedSharding`. It is primarily a helper
+    for :func:`infer_resource_partitions` but may be useful when only the partition
+    specification is required.
 
-    If use_auto_sharding is True, then we use the new experimental AUTO-sharding feature, which is not yet
-    fully supported by JAX. If it is False, then we will guess fully replicated for any unnamed arrays that
-    don't have a sharding.
+    If ``preserve_existing_shardings`` is ``True``, then arrays that already have a
+    sharding are left untouched and ``None`` is returned for those leaves.
     """
     if resource_mapping is None:
         resource_mapping = current_thread_local_mapping()
@@ -202,35 +200,57 @@ def infer_resource_partitions(
     if resource_mapping is None:
         raise ValueError("No resource mapping found")
 
-    mesh = mesh or _get_mesh()
-    assert not isinstance(mesh, dict)
-
     def partition_spec(node: typing.Any):
         if isinstance(node, NamedArray):
             # If our NamedArray doesn't have an array (or a shapedtypestruct), we can't shard it
-            # so better to not try
             if not is_jax_array_like(node.array):
                 return None
 
-            if preserve_existing_shardings:
-                current_sharding = getattr(node.array, "sharding", None)
-            else:
-                current_sharding = None
-
+            current_sharding = getattr(node.array, "sharding", None) if preserve_existing_shardings else None
             if current_sharding is not None:
-                return current_sharding
+                return None
             else:
-                sharding = NamedSharding(mesh, pspec_for_axis(node.axes, resource_mapping))
-                return sharding
+                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?
             if node.shape == ():
-                return NamedSharding(mesh, PartitionSpec())
+                return PartitionSpec()
             elif isinstance(sharding, SingleDeviceSharding):
-                return NamedSharding(mesh, PartitionSpec(None))
-            elif sharding is not None:
-                return sharding
+                return PartitionSpec(None)
+            elif sharding is not None and preserve_existing_shardings:
+                return None
             # elif use_auto_sharding:
             # TODO: auto doesn't seem to really work reliably yet
             #     compat between 0.4.10 and 0.4.11
@@ -238,13 +258,52 @@ def infer_resource_partitions(
             #     return AUTO(mesh)
             # else:
             #     return AUTO
-            return NamedSharding(mesh, PartitionSpec(None))
+            return PartitionSpec(None)
         elif isinstance(node, (bool, float, complex, int)):
-            return NamedSharding(mesh, PartitionSpec())
+            return PartitionSpec()
         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(
+    tree: PyTree,
+    resource_mapping: Optional[ResourceMapping] = None,
+    preserve_existing_shardings: bool = True,
+    use_auto_sharding: bool = True,
+    mesh: Optional[Mesh] = None,
+) -> PyTree:
+    """
+    Infer the sharding for a module, to be used with ``named_jit``.
+
+    This first calls :func:`pspec_for` to compute ``PartitionSpec`` objects and then
+    wraps them in :class:`~jax.sharding.NamedSharding` using the provided mesh. If
+    ``preserve_existing_shardings`` is ``True``, then arrays that are already sharded
+    retain their current sharding.
+    """
+    pspecs = pspec_for(
+        tree,
+        resource_mapping=resource_mapping,
+        preserve_existing_shardings=preserve_existing_shardings,
+        use_auto_sharding=use_auto_sharding,
+    )
+
+    mesh = mesh or _get_mesh()
+    assert not isinstance(mesh, dict)
+
+    def to_sharding(node: typing.Any, spec: typing.Any):
+        if spec is None:
+            if isinstance(node, NamedArray):
+                return getattr(node.array, "sharding", None)
+            elif is_jax_array_like(node):
+                return getattr(node, "sharding", None)
+            else:
+                return None
+        else:
+            return NamedSharding(mesh, spec)
+
+    return htu.tree_map(to_sharding, tree, pspecs)
 
 
 class WrappedCallable(typing.Protocol[Args, R]):
@@ -659,6 +718,7 @@ __all__ = [
     "auto_sharded",
     "shard",
     "shard_with_axis_mapping",
+    "pspec_for",
     "infer_resource_partitions",
     "named_jit",
     "fsdp",

haliax-1.4.dev403/tests/test_field.py

@@ -0,0 +1,19 @@
+import jax.numpy as jnp
+import equinox as eqx
+import pytest
+
+import haliax as hax
+
+
+class M(eqx.Module):
+    a: jnp.ndarray = hax.field(axis_names=("batch",))
+
+
+def test_axis_names_metadata():
+    field = M.__dataclass_fields__["a"]
+    assert field.metadata["axis_names"] == ("batch",)
+
+
+def test_axis_names_static_exclusive():
+    with pytest.raises(ValueError):
+        hax.field(static=True, axis_names=("x",))

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

@@ -7,7 +7,13 @@ from jaxtyping import Array
 
 import haliax as hax
 from haliax import Axis, NamedArray
-from haliax.partitioning import ResourceAxis, axis_mapping, infer_resource_partitions, named_jit
+from haliax.partitioning import (
+    ResourceAxis,
+    axis_mapping,
+    infer_resource_partitions,
+    named_jit,
+    pspec_for,
+)
 from test_utils import skip_if_not_enough_devices
 
 
@@ -40,6 +46,47 @@ def test_infer_named_axes():
         assert axes.unnamed1.is_fully_replicated
 
 
+def test_pspec_for_named_axes():
+    mesh = Mesh(np.array(jax.devices()).reshape(-1, 1), (ResourceAxis.DATA, ResourceAxis.MODEL))
+    with axis_mapping(resource_map), mesh:
+        mod = MyModule(named=hax.ones((Dim1, Dim2, Dim3)), unnamed1=jnp.ones(Dim2.size), static_field=1)
+
+        specs: MyModule = pspec_for(mod, preserve_existing_shardings=False)
+
+        spec = PartitionSpec(None, ResourceAxis.DATA, ResourceAxis.MODEL)
+
+        assert specs.named == spec
+        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.dev400/src/haliax/__about__.py

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