haliax 1.4.dev330__py3-none-any.whl → 1.4.dev332__py3-none-any.whl

haliax/__about__.py

@@ -1 +1 @@
-__version__ = "1.4.dev330"
+__version__ = "1.4.dev332"

haliax/__init__.py

@@ -61,7 +61,7 @@ from .core import (
     updated_slice,
 )
 from .hof import fold, map, scan, vmap
-from .jax_utils import filter_checkpoint
+from .jax_utils import tree_checkpoint_name
 from .ops import clip, isclose, pad_left, trace, tril, triu, where
 from .partitioning import auto_sharded, axis_mapping, fsdp, named_jit, shard, shard_with_axis_mapping
 from .specialized_fns import top_k
@@ -887,7 +887,6 @@ def true_divide(x1: NamedOrNumeric, x2: NamedOrNumeric, /) -> NamedOrNumeric:
 # deprecated name
 concat_axis_specs = concat_axes
 
-
 __all__ = [
     "debug",
     "random",
@@ -1071,4 +1070,5 @@ __all__ = [
     "ravel",
     "flatten",
     "is_named_array",
+    "tree_checkpoint_name",
 ]

haliax/jax_utils.py

@@ -1,5 +1,6 @@
 import functools as ft
 import typing
+import warnings
 from typing import Any, Callable, Optional, Sequence, Union
 
 import equinox as eqx
@@ -8,8 +9,8 @@ import numpy as np
 from jax import Array
 from jax import numpy as jnp
 from jax import random as jrandom
-from jax._src.numpy import lax_numpy
-from jax._src.typing import DTypeLike
+from jax.ad_checkpoint import checkpoint_name
+from jax.typing import DTypeLike
 from jaxtyping import PRNGKeyArray
 
 import haliax
@@ -27,6 +28,7 @@ except ImportError:
 
 
 F = typing.TypeVar("F", bound=Callable[..., Any])
+T = typing.TypeVar("T")
 
 
 class Static(eqx.Module):
@@ -70,23 +72,9 @@ def filter_eval_shape(*args, **kwargs):
 def filter_checkpoint(fun: Callable, *, prevent_cse: bool = True, policy: Optional[Callable[..., bool]] = None):
     """As `jax.checkpoint`, but allows any Python object as inputs and outputs"""
 
-    @ft.wraps(fun)
-    def _fn(_static, _dynamic):
-        _args, _kwargs = eqx.combine(_static, _dynamic)
-        _out = fun(*_args, **_kwargs)
-        _dynamic_out, _static_out = eqx.partition(_out, is_jax_array_like)
-        return _dynamic_out, Static(_static_out)
+    warnings.warn("filter_checkpoint is deprecated, use eqx.filter_checkpoint instead", DeprecationWarning)
 
-    checkpointed_fun = jax.checkpoint(_fn, prevent_cse=prevent_cse, policy=policy, static_argnums=(0,))
-
-    @ft.wraps(fun)
-    def wrapper(*args, **kwargs):
-        dynamic, static = eqx.partition((args, kwargs), is_jax_array_like)
-        dynamic_out, static_out = checkpointed_fun(static, dynamic)
-
-        return eqx.combine(dynamic_out, static_out.value)
-
-    return wrapper
+    return eqx.filter_checkpoint(fun, prevent_cse=prevent_cse, policy=policy)
 
 
 def is_jax_array_like(x):
@@ -202,7 +190,7 @@ def _jittable_dg_einsum(
         contract_path = opt_einsum.contract_path
     else:
         ty = next(iter(non_constant_dim_types))
-        contract_path = lax_numpy._poly_einsum_handlers.get(ty, lax_numpy._default_poly_einsum_handler)
+        contract_path = jax_einsum._poly_einsum_handlers.get(ty, jax_einsum._default_poly_einsum_handler)
     # using einsum_call=True here is an internal api for opt_einsum... sorry
     operands, contractions = contract_path(*operands, einsum_call=True, use_blas=True, optimize=optimize)
 
@@ -212,3 +200,23 @@ def _jittable_dg_einsum(
     if spec is not None:
         einsum = jax.named_call(einsum, name=spec)
     return einsum(operands, contractions, precision, preferred_element_type, _dot_general)  # type: ignore[operator]
+
+
+def tree_checkpoint_name(x: T, name: str) -> T:
+    """
+    Checkpoint a tree of arrays with a given name. This is useful for gradient checkpointing.
+    This is equivalent to calling [jax.ad_checkpoint.checkpoint_name][]
+    except that it works for any PyTree, not just arrays.
+
+    See Also:
+        * [jax.ad_checkpoint.checkpoint_name][]
+        * [haliax.nn.StackedCheckpointPolicy][]
+    """
+
+    def _checkpoint_leaf(x):
+        if is_jax_array_like(x):
+            return checkpoint_name(x, name)
+        else:
+            return x
+
+    return jax.tree.map(_checkpoint_leaf, x)

haliax/nn/__init__.py

@@ -38,10 +38,7 @@ from .loss import binary_cross_entropy_loss, cross_entropy_loss, cross_entropy_l
 from .mlp import MLP
 from .normalization import LayerNorm, log_softmax, logsumexp, softmax, standardize
 from .pool import max_pool, mean_pool, min_pool
-from .scan import BlockSeq, Stacked
-
-
-# TODO: support where in softmax, etc
+from .scan import BlockSeq, Stacked, StackedCheckpointPolicy
 
 
 def one_hot(x: NamedArray | int, class_axis: Axis, *, dtype=None) -> NamedArray:

haliax/nn/scan.py

@@ -1,6 +1,8 @@
+import dataclasses
 import functools
 import re
-from typing import Any, Dict, Generic, Optional, Protocol, Sequence, Type, TypeVar, cast
+import warnings
+from typing import Any, Dict, Generic, Literal, Optional, Protocol, Sequence, Type, TypeVar, Union, cast
 
 import equinox as eqx
 import jax
@@ -8,7 +10,7 @@ from jax import numpy as jnp
 
 import haliax
 import haliax.util
-from haliax.jax_utils import filter_checkpoint, is_jax_array_like
+from haliax.jax_utils import tree_checkpoint_name
 from haliax.util import is_jax_or_hax_array_like
 
 from .._src.state_dict import ModuleWithStateDictSerialization, StateDict, with_prefix
@@ -26,6 +28,179 @@ class ModuleInit(Protocol[M_co]):
         ...
 
 
+@dataclasses.dataclass
+class StackedCheckpointPolicy:
+    """
+    A class that represents a gradient checkpoint policy for blocks in a Stacked module. This is used to control
+    gradient checkpointing in [haliax.nn.Stacked][] and [haliax.nn.BlockSeq][].
+
+    Gradient checkpointing is a technique for reducing memory usage in training large models. It works by saving only a
+    subset of the forward pass and recomputing the rest in the backward pass. (By doing parts of the forward pass again)
+    JAX suggests that this usually isn't necessary when not using scan-over-layers (i.e. Stacked), so this is mostly
+    useful for Stacked modules.
+
+    A scan block takes a "carry" and some extra arguments, and returns a "carry" and an "output". The "carry" is passed
+    to the next block, and the "output" is concatenated into a final result (sort of like an RNN).
+
+    Schematically it might look like this:
+
+    ```
+          I       I       I       I
+          |       |       |       |
+    C ->  B -C->  B -C->  B -C->  B --> C
+          |       |       |       |
+          O       O       O       O
+    ```
+
+    where "C" is the carry and "O" is the output. A block will typically do some computation (e.g. a Transformer block)
+    as well, which might require saving or recomputing in the backward pass.
+
+    Imagine we save just the carries, then during the backward pass, we can recompute the outputs using the carries
+    and the inputs (and the blocks), and then compute the gradient as usual. This requires O(N) memory and O(N) time,
+    where N is the number of blocks. This is the default behavior in Haliax and works well for most models.
+
+    Alternatively, we could only save the initial and final carry. (This corresponds to
+    `StackedCheckpointPolicy(save_carries=False, save_outputs=False)` or `"recompute"`)
+    Then, during the backward pass, for each block we
+    can compute all blocks up to that point (to get its input carry) and then compute the block itself.
+    This requires O(1) memory and O(N^2) time.
+
+    Intermediate approaches exist (including O(sqrt(N)) memory and O(N) time), but we don't support them yet.
+
+    Another choice is to "offload" carries and outputs to the host, which can reduce memory usage on the device.
+    We support offloading carries and outputs to the host, but not internals.
+
+    See Also:
+        * [JAX docs on gradient checkpointing](https://docs.jax.dev/en/latest/gradient-checkpointing.html)
+    """
+
+    save_carries: bool | Literal["offload"] = True
+    """
+    Whether to save all carries in the forward pass. If True, carries are saved in the forward pass and used in the
+    backward pass. If "offload", carries are saved in the forward pass and offloaded to the host
+    """
+    save_outputs: bool | Literal["offload"] = True
+    """
+    Whether to save scan outputs in the forward pass. If True, outputs are saved in the forward pass and
+    used in the backward pass. If "offload", outputs are saved in the forward pass and offloaded to the host
+    """
+    save_block_internals: bool | list[str] = True
+    """
+    Whether to save internal state of blocks. If a list, only the listed names are saved, as
+    with [jax.checkpoint_policies.save_only_these_names][].
+
+    See Also: https://docs.jax.dev/en/latest/gradient-checkpointing.html#custom-policies-for-offload
+    """
+    prevent_cse: bool = False
+    """
+    Whether to prevent common subexpression elimination in the checkpointed function.
+    """
+
+    disable: bool = False
+    """
+    Whether to disable gradient checkpointing entirely. This is useful for debugging.
+    """
+
+    @staticmethod
+    def from_bool_or_str(remat_policy: bool | str):
+        """
+        Convert a boolean or string into a BlockCheckpointPolicy. This is useful for converting user input
+        into a BlockCheckpointPolicy.
+
+        Choices:
+            * True: save outputs, don't save block internals. This is the classic Haliax behavior.
+            * False: save everything.
+            * "offload": offload outputs to the host, don't save block internals.
+            * "recompute" or "full": don't save outputs or block internals.
+            * "save_all": save outputs and block internals. Equivalent to False
+        """
+        if remat_policy == "offload":
+            return StackedCheckpointPolicy(save_carries="offload", save_outputs="offload", save_block_internals=False)
+        elif remat_policy == "recompute" or remat_policy == "full":
+            return StackedCheckpointPolicy(save_carries=False, save_outputs=False, save_block_internals=False)
+        elif remat_policy == "save_all":
+            return StackedCheckpointPolicy(save_carries=True, save_outputs=True, save_block_internals=True)
+        elif remat_policy is True:
+            return StackedCheckpointPolicy(save_carries=True, save_outputs=True, save_block_internals=False)
+        elif remat_policy is False:
+            return StackedCheckpointPolicy(save_carries=True, save_outputs=True, save_block_internals=True)
+        else:
+            raise ValueError(f"Invalid checkpoint policy {remat_policy}")
+
+    @staticmethod
+    def _mk(remat_policy: Union[bool, str, "StackedCheckpointPolicy"]) -> "StackedCheckpointPolicy":
+        if isinstance(remat_policy, StackedCheckpointPolicy):
+            return remat_policy
+        else:
+            return StackedCheckpointPolicy.from_bool_or_str(remat_policy)
+
+    def checkpoint(self, carry_name: str, output_name: str, callable):
+        if self.disable:
+            return callable
+        policy = self._to_jax_policy(carry_name, output_name)
+        if policy is None:
+            return callable
+        else:
+            return eqx.filter_checkpoint(callable, policy=policy, prevent_cse=self.prevent_cse)
+
+    def _to_jax_policy(self, carry_name: str, output_name: str):
+        our_names_to_save = []
+        our_names_to_offload = []
+        our_names_to_remat = []
+
+        if self.save_outputs is True:
+            our_names_to_save.append(output_name)
+        elif self.save_outputs == "offload":
+            our_names_to_offload.append(output_name)
+        else:
+            assert self.save_outputs is False, f"Invalid save_outputs {self.save_outputs}"
+            our_names_to_remat.append(output_name)
+
+        if self.save_carries is True:
+            our_names_to_save.append(carry_name)
+        elif self.save_carries == "offload":
+            our_names_to_offload.append(carry_name)
+        else:
+            assert self.save_carries is False, f"Invalid save_carries {self.save_carries}"
+            our_names_to_remat.append(carry_name)
+
+        if isinstance(self.save_block_internals, Sequence):
+            our_names_to_save.extend(self.save_block_internals)
+
+        if len(our_names_to_offload) > 0:
+            if self.save_block_internals is True:
+                raise ValueError("Can't save all block internals and offload outputs. Use a list of names instead.")
+
+            return jax.checkpoint_policies.save_and_offload_only_these_names(
+                names_which_can_be_saved=our_names_to_save,
+                names_which_can_be_offloaded=our_names_to_offload,
+                offload_src="device",
+                offload_dst="pinned_host",
+            )
+        else:
+            if len(our_names_to_remat) > 0:
+                if self.save_block_internals is True:
+                    p1 = jax.checkpoint_policies.save_anything_except_these_names(*our_names_to_remat)
+                    if len(our_names_to_save) > 0:
+                        p2 = jax.checkpoint_policies.save_only_these_names(*our_names_to_save)
+                        return jax.checkpoint_policies.save_from_both_policies(p1, p2)
+                    else:
+                        return p1
+                else:
+                    return jax.checkpoint_policies.save_only_these_names(*our_names_to_save)
+            elif len(our_names_to_save) > 0:
+                p1 = jax.checkpoint_policies.save_only_these_names(*our_names_to_save)
+                if self.save_block_internals is True:
+                    p2 = jax.checkpoint_policies.save_anything_except_these_names(*our_names_to_remat)
+                    return jax.checkpoint_policies.save_from_both_policies(p1, p2)
+                else:
+                    return p1
+            elif self.save_block_internals is True:
+                return jax.checkpoint_policies.save_anything_except_these_names(*our_names_to_remat)
+            else:
+                return None
+
+
 class BlockFoldable(Protocol[M]):
     """
     A superclass for [haliax.nn.Stacked][] and [haliax.nn.BlockSeq][] that exposes the fold and scan methods, as
@@ -39,7 +214,12 @@ class BlockFoldable(Protocol[M]):
 
     @classmethod
     def init(
-        cls: Type[S], Block: Axis, module: Type[M], *, gradient_checkpointing: bool = False, prevent_cse: bool = False
+        cls: Type[S],
+        Block: Axis,
+        module: Type[M],
+        *,
+        gradient_checkpointing: bool | StackedCheckpointPolicy = False,
+        prevent_cse: bool = False,
     ) -> ModuleInit[S]:
         ...
 
@@ -70,23 +250,37 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
 
     blocks: Sequence[M]
     Block: Axis = eqx.static_field()
-    gradient_checkpointing: bool = eqx.static_field()
+    gradient_checkpointing: StackedCheckpointPolicy = eqx.static_field()
 
     @classmethod
     def init(
-        cls: Type[S], Block: Axis, module: Type[M], *, gradient_checkpointing: bool = False, prevent_cse: bool = False
+        cls: Type[S],
+        Block: Axis,
+        module: Type[M],
+        *,
+        gradient_checkpointing: bool | StackedCheckpointPolicy = False,
+        prevent_cse: bool | None = None,
     ) -> ModuleInit[S]:
         """
         This is a curried init method that takes the Block and module and returns a function that takes
         the arguments to the module's init method. Any NamedArrays in the arguments will be sliced along the
         Block axis (if it exists). JAX arrays will be sliced along the first axis.
         """
-        del prevent_cse  # not needed, but kept for compat with Stacked
+
+        gradient_checkpointing = StackedCheckpointPolicy._mk(gradient_checkpointing)
+
+        if prevent_cse is not None:
+            warnings.warn(
+                "The prevent_cse argument is deprecated and will be removed in a future version of Haliax. Use the"
+                " StackedCheckpointPolicy instead.",
+                DeprecationWarning,
+            )
+            gradient_checkpointing = dataclasses.replace(gradient_checkpointing, prevent_cse=prevent_cse)
 
         @functools.wraps(module)
         def fn(*args, **kwargs):
             # The only complexity here is that the args and kwargs might have a Block axis in them,
-            # in which case we need to loop over them them to slice them out.
+            # in which case we need to loop over them to slice them out.
 
             def init_block(i):
                 (block_args, block_kwargs) = haliax.tree_util.tree_map(
@@ -101,38 +295,50 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
         return fn
 
     def scan(self, init: T, *extra_args, **extra_kwargs):
-        out = []
-        carry = init
+        def do_scan(init, *extra_args, **extra_kwargs):
+            out = []
+            carry = init
 
-        for i, block in enumerate(self.blocks):
-            if self.gradient_checkpointing:
-                block = filter_checkpoint(block)
-            (block_args, block_kwargs) = haliax.tree_util.tree_map(
-                functools.partial(BlockSeq._slice_out, self.Block, i), (extra_args, extra_kwargs)
-            )
-            block_result = block(carry, *block_args, **block_kwargs)
-            if not isinstance(block_result, (tuple, list)) or len(block_result) != 2:
-                raise ValueError(
-                    f"BlockSeq.scan expects the block to return a pair of (carry, extra), got {block_result}"
+            for i, block in enumerate(self.blocks):
+
+                (block_args, block_kwargs) = haliax.tree_util.tree_map(
+                    functools.partial(BlockSeq._slice_out, self.Block, i), (extra_args, extra_kwargs)
                 )
 
-            carry, extra = block_result
+                block_result = block(carry, *block_args, **block_kwargs)
 
-            out.append(extra)
+                if not isinstance(block_result, (tuple, list)) or len(block_result) != 2:
+                    raise ValueError(
+                        f"BlockSeq.scan expects the block to return a pair of (carry, extra), got {block_result}"
+                    )
 
-        # TODO: do we want to stack the outputs?
-        return carry, out
+                carry, extra = block_result
+
+                carry = tree_checkpoint_name(carry, self._carry_ckpt_name)
+                extra = tree_checkpoint_name(extra, self._output_ckpt_name)
+
+                out.append(extra)
+
+            return carry, haliax.tree_util.tree_map(lambda *x: haliax.stack(self.Block, x), *out)
+
+        do_scan = self.gradient_checkpointing.checkpoint(self._carry_ckpt_name, self._output_ckpt_name, do_scan)
+
+        return do_scan(init, *extra_args, **extra_kwargs)
 
     def fold(self, init: T, *args, **kwargs) -> T:
-        carry = init
-        for i, block in enumerate(self.blocks):
-            if self.gradient_checkpointing:
-                block = filter_checkpoint(block)
-            (block_args, block_kwargs) = haliax.tree_util.tree_map(
-                functools.partial(BlockSeq._slice_out, self.Block, i), (args, kwargs)
-            )
-            carry = block(carry, *block_args, **block_kwargs)
-        return carry
+        def do_fold(init, *args, **kwargs):
+            carry = init
+            for i, block in enumerate(self.blocks):
+                (block_args, block_kwargs) = haliax.tree_util.tree_map(
+                    functools.partial(BlockSeq._slice_out, self.Block, i), (args, kwargs)
+                )
+                carry = block(carry, *block_args, **block_kwargs)
+                carry = tree_checkpoint_name(carry, self._carry_ckpt_name)
+            return carry
+
+        do_fold = self.gradient_checkpointing.checkpoint(self._carry_ckpt_name, self._output_ckpt_name, do_fold)
+
+        return do_fold(init, *args, **kwargs)
 
     def unstacked(self) -> Sequence[M]:
         return self.blocks
@@ -174,6 +380,14 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
 
         return state_dict
 
+    @property
+    def _output_ckpt_name(self):
+        return f"BlockSeq[{self.Block}, {self.blocks[0].__class__.__name__}].outputs"
+
+    @property
+    def _carry_ckpt_name(self):
+        return f"BlockSeq[{self.Block}, {self.blocks[0].__class__.__name__}].carry"
+
 
 class Stacked(ModuleWithStateDictSerialization, Generic[M]):
     """
@@ -192,12 +406,12 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
     that the function has the same control flow for every element of the stack.
 
     Stacked supports both "fold" and "scan" semantics. "fold" is the same as a for loop that accumulates a single
-    output, while "scan" is the same as a for loop that accumulates a list of intermediates as well as the final output.
+    output, while "scan" is the same as a for loop that accumulates a list of outputs as well as the final output.
 
     Stacked also supports gradient checkpointing, which is useful for very large models that don't fit in memory.
 
     Typically only one of "fold" or "scan" can be used with a given Stacked module, depending on the what the module
-    returns: if the module returns a single output, use "fold"; if the module returns a sequence of intermediates and
+    returns: if the module returns a single output, use "fold"; if the module returns a sequence of outputs and
     an output to be passed to the next layer, use "scan". More concretely, for a transformer, you would use "scan" if
     you wanted to return a kv cache (or the attention matrix) as well as the output of the transformer. If you just
     wanted the output of the transformer, you would use "fold".
@@ -227,48 +441,65 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
 
     stacked: M
     Block: Axis = eqx.static_field()
-    # TODO: support fancier gradient checkpointing
-    gradient_checkpointing: bool = eqx.static_field()
-    prevent_cse: bool = eqx.static_field()
+    gradient_checkpointing: StackedCheckpointPolicy = eqx.static_field()
 
     @classmethod
     def init(
-        cls, Block: Axis, module: Type[M], *, gradient_checkpointing: bool = False, prevent_cse: bool = False
+        cls,
+        Block: Axis,
+        module: Type[M],
+        *,
+        gradient_checkpointing: bool | StackedCheckpointPolicy | str = False,
+        prevent_cse: bool | None = None,
     ) -> ModuleInit["Stacked[M]"]:
         """
         Initialize a Stacked module. This method is curried: you can pass in the Block and module, and it will return
         a function that takes (batched) arguments to the vmapped module's init method.
-        :param Block:
-        :param module:
-        :param gradient_checkpointing:
-        :param prevent_cse:
-        :return:
+
+        Args:
+            Block: The axis that will be stacked over. This is typically a "layer" axis, but could be any axis.
+            module: The module that will be stacked. This module must take a batched input and return a batched output.
+            gradient_checkpointing: Whether to use gradient checkpointing. If True, uses the default policy. If a string,
+                uses the policy specified by the string. If a StackedCheckpointPolicy, uses that policy.
+            prevent_cse: Whether to prevent common subexpression elimination in the checkpointed function. This is useful
+                for debugging, but may slow down the function.
         """
 
+        gradient_checkpointing = StackedCheckpointPolicy._mk(gradient_checkpointing)
+
+        if prevent_cse is not None:
+            warnings.warn(
+                "The prevent_cse argument is deprecated and will be removed in a future version of Haliax. Use the"
+                " StackedCheckpointPolicy instead.",
+                DeprecationWarning,
+            )
+
+            gradient_checkpointing = dataclasses.replace(gradient_checkpointing, prevent_cse=prevent_cse)
+
         @functools.wraps(module)
         def fn(*args, **kwargs):
             stacked = haliax.vmap(module.init, Block)(*args, **kwargs)
-            return Stacked(stacked, Block, gradient_checkpointing, prevent_cse)
+            return Stacked(stacked, Block, gradient_checkpointing)
 
         return fn
 
     def scan(self, init, *extra_args, **extra_kwargs):
         """
         Scan over the stacked module. This is the same as a for loop that applies each instance of the module in sequence
-        to the input, passing the output of one instance to the next instance. It returns a stack of intermediates as
+        to the input, passing the output of one instance to the next instance. It returns a stack of outputs as
         well as the final output.
 
         That is, it behaves similarly to the following Python code:
 
         ```python
         carry = init
-        intermediates = []
+        outputs = []
 
         for block in self.stacked:
             carry, extra = block(carry)
-            intermediates.append(extra)
+            outputs.append(extra)
 
-        return carry, hax.stack(Block, intermediates)
+        return carry, hax.stack(Block, outputs)
         ```
 
         Args:
@@ -279,11 +510,24 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
         Returns:
 
         """
-        if self.gradient_checkpointing:
-            do_block = filter_checkpoint(self._do_block, prevent_cse=self.prevent_cse)
-        else:
-            do_block = self._do_block
-        return haliax.scan(do_block, self.Block)(init, self.stacked, *extra_args, **extra_kwargs)
+        carry_name = self._carry_ckpt_name
+        output_name = self._output_ckpt_name
+
+        def do_block(carry, block, *args, **kwargs):
+            carry, out = block(carry, *args, **kwargs)
+            carry = tree_checkpoint_name(carry, carry_name)
+            out = tree_checkpoint_name(out, output_name)
+            return carry, out
+
+        def do_scan(init, *extra_args, **extra_kwargs):
+            carry, out = haliax.scan(do_block, self.Block)(init, self.stacked, *extra_args, **extra_kwargs)
+            # carry = _tree_checkpoint_name(carry, carry_name)
+            # out = _tree_checkpoint_name(out, output_name)
+            return carry, out
+
+        do_scan = self.gradient_checkpointing.checkpoint(carry_name, output_name, do_scan)
+
+        return do_scan(init, *extra_args, **extra_kwargs)
 
     def fold(self, init, *args, **kwargs):
         """
@@ -300,19 +544,26 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
         ```
 
         Args:
-            init:
-            *args:
-            **kwargs:
+            init: The initial value of carry to pass to the first block
+            *args: Extra arguments to pass to the blocks. These are passed directly to the blocks
+            **kwargs: Extra keyword arguments to pass to the blocks. These are passed directly to the blocks
 
         Returns:
 
         """
-        if self.gradient_checkpointing:
-            do_block = filter_checkpoint(self._do_block, prevent_cse=self.prevent_cse)
-        else:
-            do_block = self._do_block
+        carry_name = self._carry_ckpt_name
 
-        return haliax.fold(do_block, self.Block)(init, self.stacked, *args, **kwargs)
+        def do_block(carry, block, *args, **kwargs):
+            carry = block(carry, *args, **kwargs)
+            carry = tree_checkpoint_name(carry, carry_name)
+            return carry
+
+        def do_fold(init, *extra_args, **extra_kwargs):
+            return haliax.fold(do_block, self.Block)(init, self.stacked, *extra_args, **extra_kwargs)
+
+        do_scan = self.gradient_checkpointing.checkpoint(carry_name, self._output_ckpt_name, do_fold)
+
+        return do_scan(init, *args, **kwargs)
 
     @staticmethod
     def _do_block(carry, block, *extra_args, **extra_kwargs):
@@ -363,6 +614,14 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
         out = super().from_state_dict(stacked, prefix=prefix)  # type: ignore
         return out
 
+    @property
+    def _carry_ckpt_name(self):
+        return f"Stacked[{self.Block}, {self.stacked.__class__.__name__}].carry"
+
+    @property
+    def _output_ckpt_name(self):
+        return f"Stacked[{self.Block}, {self.stacked.__class__.__name__}].outputs"
+
 
 def _stack_state_dict(state_dict: StateDict, prefix: Optional[str] = None) -> StateDict:
     """

haliax/random.py

@@ -24,7 +24,6 @@ def uniform(
     minval = broadcast_to(minval, shape).array
     maxval = broadcast_to(maxval, shape).array
     jax_shape = _to_jax_shape(shape)
-    print(jax_shape, minval, maxval)
     jax_array = jrandom.uniform(key=key, shape=jax_shape, dtype=dtype, minval=minval, maxval=maxval)
     return haliax.auto_sharded(NamedArray(jax_array, shape))
 

{haliax-1.4.dev330.dist-info → haliax-1.4.dev332.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev330
+Version: 1.4.dev332
 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/
@@ -61,8 +61,8 @@ Despite the focus on legibility, Haliax
 is also **fast**, typically about as fast as "pure" JAX code.
 Haliax is also built to be **scalable**: it
 can support [Fully-Sharded Data Parallelism (FSDP)](https://engineering.fb.com/2021/07/15/open-source/fsdp/) and Tensor Parallelism with [just a few lines of code](https://colab.research.google.com/drive/1QX4yH3zRFF3Xiibf1aahETcSQ5nbcUMz). Haliax powers [Levanter](https://github.com/stanford-crfm/levanter),
-our companion library for training large language models and other foundation models, with scale proven up to 20B parameters
-and up to a TPU v3-256 pod slice.
+our companion library for training large language models and other foundation models, with scale proven up to 70B parameters
+and up to TPU v4-2048.
 
 ## Example: Attention
 

{haliax-1.4.dev330.dist-info → haliax-1.4.dev332.dist-info}/RECORD

@@ -1,14 +1,14 @@
-haliax/__about__.py,sha256=Ti7Y_fjcyYA2er6-VH68C3VRUnF4goWC22riUkEqT58,27
-haliax/__init__.py,sha256=kQRQMPmSqkqYwxxi00b6CV3FbZWh2ShN3HMYdJOEPaE,29045
+haliax/__about__.py,sha256=1Bc7sSYtp49Iv6Rc1ZOV2dLFJ-vdQxZRlu4f9wKlkYg,27
+haliax/__init__.py,sha256=CQHrLfNXSO8hd27zqqzDeJP6Vzls_BKAE65VA48wlAc,29075
 haliax/axis.py,sha256=U8Ugd2d-u55RoQQqRiITKm1mzsZkDh3rThfAgkTWOwo,20757
 haliax/core.py,sha256=d-9nppDP_-IcxF1yHXDzzVNbP7WIadA1ojb3xHP_-SQ,71000
 haliax/debug.py,sha256=0qEGgIsEw3Jkp40oxgBAjD0Ps-cxi3QNP42uMRGlM84,3612
 haliax/hof.py,sha256=jDcii2IwAfhoNYTCYsdeS8vpSKmtzioPKZ6LBAVuWn8,18205
-haliax/jax_utils.py,sha256=fho2FYJUGbIvfg2ghU4Oll1bXEVYOI6KcZLXc9TRqUA,6959
+haliax/jax_utils.py,sha256=Kz7GeLBp7740AA0z7UBe1aFNfIPK25J8WUtJtbHWzsU,7131
 haliax/ops.py,sha256=JZOinmXbG3VmKl4HhgeH6pD3mQ83PElNPRqNIEQp3BU,5741
 haliax/partitioning.py,sha256=YzR7AJd1cy3HO0oGAJc5o0I928UD1m_OjvqMaNGccEA,23111
 haliax/quantization.py,sha256=3CkqXsMp9wXesEKCwvBrlODEmaPPHR11hIwOFb2HKdM,10591
-haliax/random.py,sha256=DsRwdV_pvXHtYxmVjuzfPTRjaCpqqNssLp7DiZwJp9k,13114
+haliax/random.py,sha256=Qasyt7CSqbcUC9-Q715dGsctpZS0qzHKJ-9smUzoP78,13077
 haliax/specialized_fns.py,sha256=reaG54m6BFTDEgInpylEaZ1YRnKlwuwLikEw7Jv3vDQ,1470
 haliax/state_dict.py,sha256=xsuxktpgzYRlc43IizwmEX0sIjnAypfGprrAugLWjWQ,1664
 haliax/tree_util.py,sha256=2S6455r1pYImF9ROHdjXVoybok_fcf18yLUtySaswuw,5582
@@ -24,7 +24,7 @@ haliax/_src/parsing.py,sha256=9JDE3unCjS8AN_B5kiAy1yV0AemYW0jL8A_T_bdcLHs,11189
 haliax/_src/rearrange.py,sha256=SWZAIMSrqQzqwZDxrMk25q6IOaK4PyDUEm-VOSe0N18,19498
 haliax/_src/state_dict.py,sha256=YcOtA_7B99_0VYBI1-WLHb4ZcxIfxQ1f1UyollTXlPI,17276
 haliax/_src/util.py,sha256=pXizGeJfRcx7QA1YVAxvru2EnMGIXAhZdcPGkwatP_M,1566
-haliax/nn/__init__.py,sha256=45z1rwMUyfsJtacjvYhxHMxZpd5jrkdxAwulHNF6uK0,2890
+haliax/nn/__init__.py,sha256=WeiO9JGBK-Am0NZ-s2UW7OmR4cFZbQAakuTDcMMnLg4,2875
 haliax/nn/activations.py,sha256=9h1uJ2pUWmb_OsRhqUGUJ6s95WZgxABDx7YWFD2tn5A,1723
 haliax/nn/attention.py,sha256=5P6IvQ5G-hpStKFnZ4mEV379Af6pkCyIA0qWXMM3Ed4,11942
 haliax/nn/conv.py,sha256=ozdvYBnb20CDz_xMKiXdpOW0ZMBZKlwadx21XafFw3E,15131
@@ -35,8 +35,8 @@ haliax/nn/loss.py,sha256=OsXpiidiKOD6gslXC54hywn9qm3shH9qyhx8baaps0Y,4565
 haliax/nn/mlp.py,sha256=KW_7phxbt0C8CLEVahP5Zoq4SQruiGyg7ltRvLrTdPg,3695
 haliax/nn/normalization.py,sha256=4e1rpVvZKW9w5eUsn8j0tBF6lAIs5YljsvOE_HaoOD8,2652
 haliax/nn/pool.py,sha256=DHeswGramsqNvYufl36NbwgWgI-Iht9bnJy170Pwfu0,8530
-haliax/nn/scan.py,sha256=0AENdvQXGV5c4h2jvbuHKMKF2eZVj2U4nZeTtldCZl4,16983
-haliax-1.4.dev330.dist-info/METADATA,sha256=A1aFm9W-6mQvL7NKrh17eRfCB0Y8RHXOPIc2mvlhHqc,7674
-haliax-1.4.dev330.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haliax-1.4.dev330.dist-info/licenses/LICENSE,sha256=bJiay7Nn5SHQ2n_4ZIT3AE0W1RGq4O7pxOApgBsaT64,11349
-haliax-1.4.dev330.dist-info/RECORD,,
+haliax/nn/scan.py,sha256=cm5GHVjdv8nWKKlo9hHzlyoj0Ubq9770EgqsdbUc7Z0,29163
+haliax-1.4.dev332.dist-info/METADATA,sha256=Qcs512OIq18GGpSwTgAhbstqbxdVfhfm5RjuueAYPoA,7663
+haliax-1.4.dev332.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+haliax-1.4.dev332.dist-info/licenses/LICENSE,sha256=bJiay7Nn5SHQ2n_4ZIT3AE0W1RGq4O7pxOApgBsaT64,11349
+haliax-1.4.dev332.dist-info/RECORD,,