haliax 1.4.dev399__tar.gz → 1.4.dev402__tar.gz

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

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

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

{haliax-1.4.dev399 → haliax-1.4.dev402}/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.dev399 → haliax-1.4.dev402}/src/haliax/core.py

@@ -15,7 +15,7 @@ import numpy as np
 
 import haliax
 import haliax.axis
-from haliax.jax_utils import is_jax_array_like, is_pallas_dslice
+from haliax.jax_utils import ensure_scalar, is_jax_array_like, is_pallas_dslice
 from haliax.util import ensure_tuple
 
 from ._src.util import index_where, py_slice, slice_t
@@ -1115,9 +1115,7 @@ def updated_slice(
         if axis_index is None:
             raise ValueError(f"axis {axis} not found in {array}")
         if isinstance(s, NamedArray):  # this can happen in the vmap case
-            if s.ndim != 0:
-                raise ValueError(f"NamedArray {s} must be a scalar for axis {axis} in updated_slice")
-            s = s.scalar()
+            s = ensure_scalar(s, name=str(axis))
 
         array_slice_indices[axis_index] = s
         total_length = array.axes[axis_index].size
@@ -1376,10 +1374,7 @@ def roll(
     if axis_indices is None:
         raise ValueError(f"axis {axis} not found in {array}")
 
-    if isinstance(shift, NamedArray):
-        if shift.ndim != 0:
-            raise TypeError("shift must be a scalar NamedArray")
-        shift = shift.array
+    shift = ensure_scalar(shift, name="shift")
 
     return NamedArray(jnp.roll(array.array, shift, axis_indices), array.axes)
 

haliax-1.4.dev402/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.dev399 → haliax-1.4.dev402}/src/haliax/jax_utils.py

@@ -153,6 +153,21 @@ def is_scalarish(x):
         return jnp.isscalar(x) or (getattr(x, "shape", None) == ())
 
 
+def ensure_scalar(x, *, name: str = "value"):
+    """Return ``x`` if it is not a :class:`NamedArray`, otherwise ensure it is a scalar.
+
+    This is useful for APIs that can accept either Python scalars or scalar
+    ``NamedArray`` objects (for example ``roll`` or ``updated_slice``).  If ``x``
+    is a ``NamedArray`` with rank greater than 0 a :class:`TypeError` is raised.
+    """
+
+    if isinstance(x, haliax.NamedArray):
+        if x.ndim != 0:
+            raise TypeError(f"{name} must be a scalar NamedArray")
+        return x.array
+    return x
+
+
 def is_on_mac_metal():
     return jax.devices()[0].platform.lower() == "metal"
 

{haliax-1.4.dev399 → haliax-1.4.dev402}/src/haliax/ops.py

@@ -9,7 +9,7 @@ import haliax
 
 from .axis import Axis, AxisSelector, axis_name
 from .core import NamedArray, NamedOrNumeric, broadcast_arrays, broadcast_arrays_and_return_axes, named
-from .jax_utils import is_scalarish
+from .jax_utils import ensure_scalar, is_scalarish
 
 
 def trace(array: NamedArray, axis1: AxisSelector, axis2: AxisSelector, offset=0, dtype=None) -> NamedArray:
@@ -89,7 +89,7 @@ def where(
             x = named(x, ())
         x, y = broadcast_arrays(x, y)
         if isinstance(condition, NamedArray):
-            condition = condition.scalar()
+            condition = ensure_scalar(condition, name="condition")
         return jax.lax.cond(condition, lambda _: x, lambda _: y, None)
 
     condition, x, y = broadcast_arrays(condition, x, y)  # type: ignore

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

@@ -177,24 +177,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 +199,26 @@ 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 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,15 +226,54 @@ 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)
 
 
+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]):
     """
     A wrapper for a callable that preserves the original function's name and qualname.
@@ -659,6 +686,7 @@ __all__ = [
     "auto_sharded",
     "shard",
     "shard_with_axis_mapping",
+    "pspec_for",
     "infer_resource_partitions",
     "named_jit",
     "fsdp",

{haliax-1.4.dev399 → haliax-1.4.dev402}/src/haliax/wrap.py

@@ -5,7 +5,7 @@ import jax
 from haliax.core import NamedArray, _broadcast_order, broadcast_to
 
 from .axis import AxisSelection, AxisSelector, axis_spec_to_shape_dict, eliminate_axes
-from .jax_utils import is_scalarish
+from .jax_utils import ensure_scalar, is_scalarish
 
 
 def wrap_elemwise_unary(f, a, *args, **kwargs):
@@ -105,7 +105,7 @@ def wrap_elemwise_binary(op):
             else:
                 if is_scalarish(b):
                     return NamedArray(op(a.array, b), a.axes)
-                a = a.scalar()
+                a = ensure_scalar(a)
                 return op(a, b)
 
             return NamedArray(op(a.array, b), a.axes)
@@ -119,7 +119,7 @@ def wrap_elemwise_binary(op):
             else:
                 if is_scalarish(a):
                     return NamedArray(op(a, b.array), b.axes)
-                b = b.scalar()
+                b = ensure_scalar(b)
                 return op(a, b)
 
             return NamedArray(op(a, b.array), b.axes)

haliax-1.4.dev402/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.dev399 → haliax-1.4.dev402}/tests/test_ops.py

@@ -437,3 +437,14 @@ def test_roll_scalar_named_shift():
 
     assert rolled.axes == arr.axes
     assert jnp.all(rolled.array == expected)
+
+
+def test_roll_bad_named_shift():
+    H = Axis("H", 4)
+    W = Axis("W", 3)
+
+    arr = hax.arange((H, W))
+    shift = hax.arange((Axis("dummy", 2),))
+
+    with pytest.raises(TypeError):
+        hax.roll(arr, shift, H)

{haliax-1.4.dev399 → haliax-1.4.dev402}/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,19 @@ 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 MyModuleInit(eqx.Module):
     named: NamedArray
     unnamed1: Array

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

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