haliax 1.4.dev444__tar.gz → 1.4.dev446__tar.gz

{haliax-1.4.dev444 → haliax-1.4.dev446}/PKG-INFO

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

@@ -3,4 +3,4 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-__version__ = "1.4.dev444"
+__version__ = "1.4.dev446"

{haliax-1.4.dev444 → haliax-1.4.dev446}/src/haliax/_src/state_dict.py

@@ -13,7 +13,7 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 from jax import ShapeDtypeStruct
-from jax.experimental.multihost_utils import sync_global_devices
+from haliax.jax_utils import sync_global_devices
 from jax.sharding import Mesh, NamedSharding, PartitionSpec
 from jax.tree_util import DictKey, FlattenedIndexKey, GetAttrKey, SequenceKey
 from jaxtyping import PyTree
@@ -68,7 +68,11 @@ def flatten_modules_for_export(t: T) -> T:
             )
         return module
 
-    return scan_aware_tree_map(_flatten_module, t, is_leaf=lambda x: isinstance(x, ModuleWithStateDictSerialization))
+    return scan_aware_tree_map(
+        _flatten_module,
+        t,
+        is_leaf=lambda x: isinstance(x, ModuleWithStateDictSerialization),
+    )
 
 
 def unflatten_modules_from_export(t: T, template: T) -> T:
@@ -88,7 +92,10 @@ def unflatten_modules_from_export(t: T, template: T) -> T:
         return module
 
     return scan_aware_tree_map(
-        _unflatten_module, t, template, is_leaf=lambda x: isinstance(x, ModuleWithStateDictSerialization)
+        _unflatten_module,
+        t,
+        template,
+        is_leaf=lambda x: isinstance(x, ModuleWithStateDictSerialization),
     )
 
 
@@ -187,9 +194,13 @@ def from_state_dict(tree: T, state_dict: StateDict, prefix: str | None = None) -
         else:
             return default_eqx_module_from_state_dict(tree, state_dict, prefix)
     elif isinstance(tree, list):
-        return [from_state_dict(item, state_dict, with_prefix(prefix, str(i))) for i, item in enumerate(tree)]  # type: ignore
+        return [
+            from_state_dict(item, state_dict, with_prefix(prefix, str(i))) for i, item in enumerate(tree)
+        ]  # type: ignore
     elif isinstance(tree, dict):
-        return {k: from_state_dict(v, state_dict, prefix=with_prefix(prefix, k)) for k, v in tree.items()}  # type: ignore
+        return {
+            k: from_state_dict(v, state_dict, prefix=with_prefix(prefix, k)) for k, v in tree.items()
+        }  # type: ignore
     elif isinstance(tree, NamedArray):
         if prefix is None:
             raise ValueError("Cannot extract a leaf value from a torch dict without a prefix")
@@ -383,7 +394,8 @@ def to_numpy_state_dict(model, prefix: str | None = None) -> StateDict:
                     # TODO: ensure that this mesh arranges devices correctly
                     # (jax seems to do this internally itself, so we should be fine?)
                     process_mesh = Mesh(
-                        np.array(jax.devices()).reshape((jax.process_count(), -1)), ("process", "device")
+                        np.array(jax.devices()).reshape((jax.process_count(), -1)),
+                        ("process", "device"),
                     )
 
                     # now we need to find an axis along which we can shard the array.
@@ -391,7 +403,8 @@ def to_numpy_state_dict(model, prefix: str | None = None) -> StateDict:
 
                     try:
                         axis_to_shard = index_where(
-                            lambda axis_size: axis_size % process_mesh.devices.size == 0, arr.shape
+                            lambda axis_size: axis_size % process_mesh.devices.size == 0,
+                            arr.shape,
                         )
                     except ValueError:
                         return np.array(arr)

{haliax-1.4.dev444 → haliax-1.4.dev446}/src/haliax/jax_utils.py

@@ -6,6 +6,7 @@
 import functools as ft
 import typing
 import warnings
+import zlib
 from typing import Any, Callable, Sequence
 
 import equinox as eqx
@@ -14,6 +15,8 @@ import numpy as np
 from jax import Array
 from jax import numpy as jnp
 from jax import random as jrandom
+from jax.experimental.multihost_utils import host_local_array_to_global_array
+from jax.sharding import PartitionSpec
 from jax.ad_checkpoint import checkpoint_name
 from jax.typing import DTypeLike
 from jaxtyping import PRNGKeyArray
@@ -39,6 +42,71 @@ class Static(eqx.Module):
     value: Any = eqx.field(static=True)
 
 
+# Non-busted version of broadcast_one_to_all from jax.multihost_utils. (The issue is that  if you use a non-contiguous
+# mesh, their utility blows up because it makes a contiguous mesh.)
+
+
+def _psum(xs: Any) -> Any:
+    return jax.tree.map(lambda x: jnp.sum(x, dtype=x.dtype, axis=0), xs)
+
+
+def broadcast_one_to_all(in_tree: Any, is_source: bool | None = None) -> Any:
+    """Broadcast data from a source host (host 0 by default) to all other hosts.
+
+    Args:
+      in_tree: pytree of arrays - each array *must* have the same shape across the
+        hosts.
+      is_source: optional bool denoting whether the caller is the source. Only
+        'source host' will contribute the data for the broadcast. If None, then
+        host 0 is used.
+
+    Returns:
+      A pytree matching in_tree where the leaves now all contain the data from the
+      first host.
+    """
+    if jax.process_count() == 1:
+        return jax.tree.map(np.asarray, in_tree)
+
+    if is_source is None:
+        is_source = jax.process_index() == 0
+
+    devices: np.ndarray = np.array(jax.devices()).reshape(jax.process_count(), jax.local_device_count())
+    global_mesh = jax.sharding.Mesh(devices, ("processes", "local_devices"))
+    pspec = PartitionSpec("processes")
+
+    def pre_jit(x):
+        if is_source:
+            inp = x
+        else:
+            inp = np.zeros_like(x)
+        inp = np.expand_dims(inp, axis=0)
+        return host_local_array_to_global_array(inp, global_mesh, pspec)
+
+    def post_jit(x):
+        return jax.device_get(x.addressable_data(0))
+
+    with haliax.partitioning.set_mesh(global_mesh):
+        in_tree = jax.tree.map(pre_jit, in_tree)
+        out_tree = jax.jit(
+            _psum,
+            out_shardings=jax.sharding.NamedSharding(global_mesh, PartitionSpec()),
+        )(in_tree)
+        return jax.tree.map(post_jit, out_tree)
+
+
+def assert_equal(in_tree, fail_message: str = ""):
+    """Verifies that all the hosts have the same tree of values."""
+    expected = broadcast_one_to_all(in_tree)
+    if not jax.tree_util.tree_all(jax.tree_util.tree_map(lambda *x: np.all(np.equal(*x)), in_tree, expected)):
+        raise AssertionError(f"{fail_message} Expected: {expected}; got: {in_tree}.")
+
+
+def sync_global_devices(name: str):
+    """Creates a barrier across all hosts/devices."""
+    h = np.uint32(zlib.crc32(name.encode()))
+    assert_equal(h, f"sync_global_devices name mismatch ('{name}')")
+
+
 def shaped_rng_split(key, split_shape: int | Sequence[int] = 2) -> PRNGKeyArray:
     if isinstance(split_shape, int):
         num_splits = split_shape
@@ -69,14 +137,25 @@ def maybe_rng_split(key: PRNGKeyArray | None, num: int = 2):
 def filter_eval_shape(*args, **kwargs):
     import warnings
 
-    warnings.warn("filter_eval_shape is deprecated, use eqx.filter_eval_shape instead", DeprecationWarning)
+    warnings.warn(
+        "filter_eval_shape is deprecated, use eqx.filter_eval_shape instead",
+        DeprecationWarning,
+    )
     return eqx.filter_eval_shape(*args, **kwargs)
 
 
-def filter_checkpoint(fun: Callable, *, prevent_cse: bool = True, policy: Callable[..., bool] | None = None):
+def filter_checkpoint(
+    fun: Callable,
+    *,
+    prevent_cse: bool = True,
+    policy: Callable[..., bool] | None = None,
+):
     """As `jax.checkpoint`, but allows any Python object as inputs and outputs"""
 
-    warnings.warn("filter_checkpoint is deprecated, use eqx.filter_checkpoint instead", DeprecationWarning)
+    warnings.warn(
+        "filter_checkpoint is deprecated, use eqx.filter_checkpoint instead",
+        DeprecationWarning,
+    )
 
     return eqx.filter_checkpoint(fun, prevent_cse=prevent_cse, policy=policy)
 

{haliax-1.4.dev444 → haliax-1.4.dev446}/src/haliax/nn/scan.py

@@ -150,6 +150,11 @@ class BlockFoldable(Protocol[M]):
         """
         ...
 
+    def get_layer(self, index: int) -> M:
+        """Return the ``index``th layer of the folded module."""
+
+        ...
+
 
 class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
     """
@@ -377,6 +382,11 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
 
         return state_dict
 
+    def get_layer(self, index: int) -> M:
+        """Return the ``index``th block in this sequential container."""
+
+        return self.blocks[index]
+
     @property
     def _output_ckpt_name(self):
         return f"BlockSeq[{self.Block}, {self.blocks[0].__class__.__name__}].outputs"
@@ -442,6 +452,12 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
     Block: Axis = eqx.field(static=True)
     gradient_checkpointing: ScanCheckpointPolicy = eqx.field(static=True)
 
+    @property
+    def Layers(self) -> Axis:
+        """Alias for :attr:`Block` used by some downstream code."""
+
+        return self.Block
+
     @classmethod
     def init(
         cls,
@@ -700,6 +716,23 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
         unstacked_leaves = tuple(zip(*unstacked_leaves))
         return tuple(map(lambda x: jax.tree_util.tree_unflatten(structure, x), unstacked_leaves))
 
+    def get_layer(self, index: int) -> M:
+        """Return the ``index``th layer of this stacked module."""
+
+        def select_leaf(leaf):
+            if isinstance(leaf, haliax.NamedArray):
+                if haliax.selects_axis(leaf.axes, self.Block):
+                    return leaf[self.Block, index]
+                else:
+                    return leaf
+            if is_jax_or_hax_array_like(leaf):
+                if getattr(leaf, "shape", ()) and leaf.shape[0] == self.Block.size:
+                    return leaf[index]
+                return leaf
+            return leaf
+
+        return haliax.tree_util.tree_map(select_leaf, self.stacked)
+
     def to_state_dict(self, prefix: str | None = None) -> StateDict:
         # this method needs to "devectorize" the blocks, so that we have a list of blocks h.0.FOO, h.1.FOO, etc.
         # first just do the normal thing with our own dict, which we'll post-process

{haliax-1.4.dev444 → haliax-1.4.dev446}/tests/test_scan.py

@@ -51,6 +51,65 @@ def test_unstacked():
         assert hax.all(module.array == m.stacked.array[i])
 
 
+def test_get_layer_stacked():
+    class Module(eqx.Module):
+        named: hax.NamedArray
+        array: jax.Array
+        static: int = eqx.static_field()
+
+        def __call__(self, x, *, key):  # pragma: no cover - unused in this test
+            return x + self.array + self.static
+
+        @staticmethod
+        def init(named, array, static):
+            return Module(named=named, array=array, static=static)
+
+    Block = hax.Axis("block", 3)
+    E = hax.Axis("E", 4)
+
+    initial_named = hax.random.uniform(jax.random.PRNGKey(0), (Block, E))
+    arrays = jax.numpy.arange(Block.size)
+
+    stacked = Stacked.init(Block, Module)(named=initial_named, array=arrays, static=2)
+
+    layer = stacked.get_layer(1)
+
+    assert isinstance(layer, Module)
+    assert layer.static == 2
+    assert layer.named.axes == (E,)
+    assert hax.all(layer.named == initial_named["block", 1])
+    assert hax.all(layer.array == arrays[1])
+
+
+def test_get_layer_blockseq():
+    class Module(eqx.Module):
+        named: hax.NamedArray
+        array: jax.Array
+        static: int = eqx.static_field()
+
+        def __call__(self, x, *, key):  # pragma: no cover - unused in this test
+            return x + self.array + self.static
+
+        @staticmethod
+        def init(named, array, static):
+            return Module(named=named, array=array, static=static)
+
+    Block = hax.Axis("block", 3)
+    E = hax.Axis("E", 4)
+
+    initial_named = hax.random.uniform(jax.random.PRNGKey(0), (Block, E))
+    arrays = jax.numpy.arange(Block.size)
+
+    seq = BlockSeq.init(Block, Module)(named=initial_named, array=arrays, static=2)
+
+    layer = seq.get_layer(2)
+
+    assert isinstance(layer, Module)
+    assert layer.static == 2
+    assert hax.all(layer.named == initial_named["block", 2])
+    assert hax.all(layer.array == arrays[2])
+
+
 def test_vmap():
     class Module(eqx.Module):
         weight: hax.NamedArray