haliax 1.4.dev336__tar.gz → 1.4.dev339__tar.gz

{haliax-1.4.dev336 → haliax-1.4.dev339}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev336
+Version: 1.4.dev339
 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.dev336 → haliax-1.4.dev339}/mkdocs.yml

@@ -91,7 +91,6 @@ nav:
     - Rearrange: 'rearrange.md'
     - Matrix Multiplication: 'matmul.md'
   - Neural Networks: 'nn.md'
-  - Module Stacks: 'stacked.md'
   - Partitioning: 'partitioning.md'
   - Higher Order Functions: 'hof.md'
   - FP8: 'fp8.md'

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

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

{haliax-1.4.dev336 → haliax-1.4.dev339}/src/haliax/__init__.py

@@ -61,7 +61,7 @@ from .core import (
     updated_slice,
 )
 from .hof import fold, map, scan, vmap
-from .jax_utils import tree_checkpoint_name
+from .jax_utils import filter_checkpoint
 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,6 +887,7 @@ def true_divide(x1: NamedOrNumeric, x2: NamedOrNumeric, /) -> NamedOrNumeric:
 # deprecated name
 concat_axis_specs = concat_axes
 
+
 __all__ = [
     "debug",
     "random",
@@ -1070,5 +1071,4 @@ __all__ = [
     "ravel",
     "flatten",
     "is_named_array",
-    "tree_checkpoint_name",
 ]

{haliax-1.4.dev336 → haliax-1.4.dev339}/src/haliax/jax_utils.py

@@ -1,6 +1,5 @@
 import functools as ft
 import typing
-import warnings
 from typing import Any, Callable, Optional, Sequence, Union
 
 import equinox as eqx
@@ -9,8 +8,8 @@ import numpy as np
 from jax import Array
 from jax import numpy as jnp
 from jax import random as jrandom
-from jax.ad_checkpoint import checkpoint_name
-from jax.typing import DTypeLike
+from jax._src.numpy import lax_numpy
+from jax._src.typing import DTypeLike
 from jaxtyping import PRNGKeyArray
 
 import haliax
@@ -28,7 +27,6 @@ except ImportError:
 
 
 F = typing.TypeVar("F", bound=Callable[..., Any])
-T = typing.TypeVar("T")
 
 
 class Static(eqx.Module):
@@ -72,9 +70,23 @@ 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"""
 
-    warnings.warn("filter_checkpoint is deprecated, use eqx.filter_checkpoint instead", DeprecationWarning)
+    @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)
 
-    return eqx.filter_checkpoint(fun, prevent_cse=prevent_cse, policy=policy)
+    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
 
 
 def is_jax_array_like(x):
@@ -190,7 +202,7 @@ def _jittable_dg_einsum(
         contract_path = opt_einsum.contract_path
     else:
         ty = next(iter(non_constant_dim_types))
-        contract_path = jax_einsum._poly_einsum_handlers.get(ty, jax_einsum._default_poly_einsum_handler)
+        contract_path = lax_numpy._poly_einsum_handlers.get(ty, lax_numpy._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)
 
@@ -200,23 +212,3 @@ 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-1.4.dev336 → haliax-1.4.dev339}/src/haliax/nn/__init__.py

@@ -38,7 +38,10 @@ 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, StackedCheckpointPolicy
+from .scan import BlockSeq, Stacked
+
+
+# TODO: support where in softmax, etc
 
 
 def one_hot(x: NamedArray | int, class_axis: Axis, *, dtype=None) -> NamedArray:

haliax-1.4.dev339/src/haliax/nn/scan.py

@@ -0,0 +1,425 @@
+import functools
+import re
+from typing import Any, Dict, Generic, Optional, Protocol, Sequence, Type, TypeVar, cast
+
+import equinox as eqx
+import jax
+from jax import numpy as jnp
+
+import haliax
+import haliax.util
+from haliax.jax_utils import filter_checkpoint, is_jax_array_like
+from haliax.util import is_jax_or_hax_array_like
+
+from .._src.state_dict import ModuleWithStateDictSerialization, StateDict, with_prefix
+from ..axis import Axis
+
+
+M = TypeVar("M", bound=eqx.Module)
+M_co = TypeVar("M_co", bound=eqx.Module, covariant=True)
+S = TypeVar("S", bound=eqx.Module)
+T = TypeVar("T")
+
+
+class ModuleInit(Protocol[M_co]):
+    def __call__(self, *args, **kwargs) -> M_co:
+        ...
+
+
+class BlockFoldable(Protocol[M]):
+    """
+    A superclass for [haliax.nn.Stacked][] and [haliax.nn.BlockSeq][] that exposes the fold and scan methods, as
+    well as a few other methods that are useful for these modules.
+
+    This is a protocol, so you can use it as a type hint for a function that takes a Stacked or BlockSeq.
+    Equinox modules can't directly inherit from Protocols, but you can use it as a type hint.
+    """
+
+    Block: Axis
+
+    @classmethod
+    def init(
+        cls: Type[S], Block: Axis, module: Type[M], *, gradient_checkpointing: bool = False, prevent_cse: bool = False
+    ) -> ModuleInit[S]:
+        ...
+
+    def scan(self, init: T, *extra_args, **extra_kwargs):
+        ...
+
+    def fold(self, init: T, *args, **kwargs) -> T:
+        ...
+
+    def unstacked(self) -> Sequence[M]:
+        """
+        Returns the unstacked version of this module. This is useful for logging or saving checkpoints.
+
+        """
+        ...
+
+
+class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
+    """
+    A "BlockSeq" wraps another module and produces a "sequential" version of it, where an input is applied
+    to each instance of the sequential module in sequence. This is useful for e.g. transformers
+    where you have multiple instances of the same transformer block and the input is applied in a fold/for loop
+    in sequence.
+
+    It's similar in spirit to an [equinox.nn.Sequential][]. Unlike [equinox.nn.Sequential][], BlockSeq does not need to be
+    homogeneous (though the init method assumes that it is).
+    """
+
+    blocks: Sequence[M]
+    Block: Axis = eqx.static_field()
+    gradient_checkpointing: bool = eqx.static_field()
+
+    @classmethod
+    def init(
+        cls: Type[S], Block: Axis, module: Type[M], *, gradient_checkpointing: bool = False, prevent_cse: bool = False
+    ) -> 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
+
+        @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.
+
+            def init_block(i):
+                (block_args, block_kwargs) = haliax.tree_util.tree_map(
+                    functools.partial(BlockSeq._slice_out, Block, i), (args, kwargs)
+                )
+                return module.init(*block_args, **block_kwargs)
+
+            seq = [init_block(i) for i in range(Block.size)]
+
+            return BlockSeq(seq, Block, gradient_checkpointing)
+
+        return fn
+
+    def scan(self, init: T, *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}"
+                )
+
+            carry, extra = block_result
+
+            out.append(extra)
+
+        # TODO: do we want to stack the outputs?
+        return carry, out
+
+    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 unstacked(self) -> Sequence[M]:
+        return self.blocks
+
+    @staticmethod
+    def _slice_out(Block, i, x):
+        if haliax.is_named_array(x):
+            if haliax.selects_axis(x.axes, Block):
+                return x[Block, i]
+            else:
+                return x
+        elif haliax.jax_utils.is_jax_array_like(x):
+            return x[i]
+        else:
+            return x
+
+    def _state_dict_key_map(self) -> Dict[str, Optional[str]]:
+        return {"blocks": None}
+
+    def from_state_dict(self: M, state_dict: StateDict, prefix: Optional[str] = None) -> M:
+        out_blocks = []
+        for i, block in enumerate(self.blocks):
+            my_prefix = with_prefix(prefix, str(i))
+            block = block.from_state_dict(state_dict, my_prefix)
+            out_blocks.append(block)
+
+        return eqx.tree_at(lambda m: m.blocks, self, out_blocks)
+
+    def to_state_dict(self, prefix: Optional[str] = None) -> StateDict:
+        """
+        Returns the unstacked format of the module, which is compatible with torch.nn.Sequential, with keys of the form (...). The stacked/vectorized format is required for haliax.nn.Stacked and vectorizes all such tensors into a single shared key.".
+        """
+        state_dict: StateDict = {}
+        for i, block in enumerate(self.blocks):
+            my_prefix = with_prefix(prefix, str(i))
+            # we can't assume to_state_dict is implemented, so we have to do it manually
+            block_dict = haliax.state_dict.to_state_dict(block, my_prefix)
+            state_dict.update(block_dict)
+
+        return state_dict
+
+
+class Stacked(ModuleWithStateDictSerialization, Generic[M]):
+    """
+    A "Stacked" wraps another module and produces a "stacked" version of it, where an input is applied
+    to each instance of the stacked module in sequence. This is useful for e.g. transformers
+    where you have multiple instances of the same transformer block and the input is applied in a fold/for loop
+    in sequence.
+
+    It's similar in spirit to an [equinox.nn.Sequential], but it must be homogeneous. In Jax, this is much cheaper to
+    compile than a sequential (or moral equivalent), because Jax compiles the module's method once, instead of unrolling
+    the sequential and compiling everything as a giant graph. In Jax, this pattern is often called "scan layers" or
+    "scan over layers".
+
+    A further constraint is that the elements of the stack must have the same Python control flow. This is because
+    Jax's scan primitive requires that the function you pass to it is pure, and the only way to do that is to ensure
+    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.
+
+    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
+    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".
+
+
+    Example:
+        ```python
+        >>> import equinox as eqx
+        >>> import haliax as hax
+        >>> import haliax.nn as hnn
+        >>> class MyModule(eqx.Module):
+        ...     def __init__(self, num_layers: int, hidden: hax.Axis, *, key):
+        ...         self.axis = hax.Axis("layer", num_layers)
+        ...         split_key = jax.random.split(key, num_layers)
+        ...         self.layers = Stacked.init(self.axis, hnn.Linear)(In=hidden, Out=hidden, key=split_key)
+        ...
+        ...     def __call__(self, x):
+        ...         return self.layers.fold(x)  # applies each layer in sequence
+        ...
+        >>> Hidden = hax.Axis("hidden", 10)
+        >>> mod = MyModule(5, Hidden)
+        >>> mod(hax.ones(Hidden))
+        ```
+    """
+
+    # TODO: we can probably make this module support pipeline parallelism, but that's a whole project in itself
+
+    stacked: M
+    Block: Axis = eqx.static_field()
+    # TODO: support fancier gradient checkpointing
+    gradient_checkpointing: bool = eqx.static_field()
+    prevent_cse: bool = eqx.static_field()
+
+    @classmethod
+    def init(
+        cls, Block: Axis, module: Type[M], *, gradient_checkpointing: bool = False, prevent_cse: bool = False
+    ) -> 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:
+        """
+
+        @functools.wraps(module)
+        def fn(*args, **kwargs):
+            stacked = haliax.vmap(module.init, Block)(*args, **kwargs)
+            return Stacked(stacked, Block, gradient_checkpointing, prevent_cse)
+
+        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
+        well as the final output.
+
+        That is, it behaves similarly to the following Python code:
+
+        ```python
+        carry = init
+        intermediates = []
+
+        for block in self.stacked:
+            carry, extra = block(carry)
+            intermediates.append(extra)
+
+        return carry, hax.stack(Block, intermediates)
+        ```
+
+        Args:
+            init:
+            *extra_args:
+            **extra_kwargs:
+
+        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)
+
+    def fold(self, init, *args, **kwargs):
+        """
+        Fold 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.
+        That is, it behaves similarly to the following Python code:
+
+        ```python
+        carry = init
+        for block in self.stacked:
+            carry = block(carry)
+
+        return carry
+        ```
+
+        Args:
+            init:
+            *args:
+            **kwargs:
+
+        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.fold(do_block, self.Block)(init, self.stacked, *args, **kwargs)
+
+    @staticmethod
+    def _do_block(carry, block, *extra_args, **extra_kwargs):
+        return block(carry, *extra_args, **extra_kwargs)
+
+    # TODO: this is for logic that's in levanter. We should move that logic to haliax I guess?
+    def _state_dict_key_map(self) -> Dict[str, Optional[str]]:
+        return {"stacked": None}
+
+    def unstacked(self) -> Sequence[M]:
+        """
+        Returns the unstacked version of this module. This is useful for logging or saving checkpoints.
+        Returns:
+            A sequence of modules, one for each element of the stack
+        """
+
+        def unbatch_leaf(x):
+            if isinstance(x, haliax.core.NamedArray):
+                if haliax.selects_axis(x.axes, self.Block):
+                    return haliax.unbind(x, self.Block)
+                else:
+                    return tuple(x for _ in range(self.Block.size))
+            elif haliax.jax_utils.is_jax_array_like(x):
+                assert (
+                    x.shape[0] == self.Block.size
+                ), f"Expected first dimension to be {self.Block.size}, got {x.shape[0]}"
+                return tuple(x[i] for i in range(self.Block.size))
+            else:
+                return tuple(x for _ in range(self.Block.size))
+
+        leaves, structure = jax.tree_util.tree_flatten(self.stacked, is_leaf=haliax.is_named_array)
+        unstacked_leaves = tuple(map(unbatch_leaf, leaves))
+        # now we need to transpose the leaves
+        unstacked_leaves = tuple(zip(*unstacked_leaves))
+        return tuple(map(lambda x: jax.tree_util.tree_unflatten(structure, x), unstacked_leaves))
+
+    def to_state_dict(self, prefix: Optional[str] = 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
+        state_dict: StateDict = super().to_state_dict(prefix)
+
+        return _unstack_state_dict(state_dict, prefix)
+
+    def from_state_dict(self: M, state_dict: StateDict, prefix: Optional[str] = None) -> M:
+        # this method needs to "vectorize" the blocks, so that we have a single block h.FOO
+        # first just do the normal thing with our own dict, which we'll post-process
+        stacked = _stack_state_dict(state_dict, prefix=prefix)
+        out = super().from_state_dict(stacked, prefix=prefix)  # type: ignore
+        return out
+
+
+def _stack_state_dict(state_dict: StateDict, prefix: Optional[str] = None) -> StateDict:
+    """
+    Stack all keys matching prefix in a new state dict, returning a state dict that has all keys matching
+    prefix stacked, but otherwise the same.
+
+    Stacked in this case means roughly "compatible with a torch.nn.Sequential", which means that the
+    keys are of the form "<prefix>.0.<key>", "<prefix>.1.<key>", etc.
+
+    Mostly for use with [haliax.nn.Stacked][].
+    """
+    vectorized_dict: StateDict = {}
+
+    tensors_to_vectorize: dict[str, list[Optional[Any]]] = {}
+    if prefix is not None:
+        prefix_for_pat = re.escape(prefix + ".")
+    else:
+        prefix_for_pat = ""
+    pattern = re.compile(rf"{prefix_for_pat}(\d+)\.(.*)")
+
+    for k, v in state_dict.items():
+        match = pattern.match(k)
+        if match:
+            block_idx = int(match.group(1))
+            block_key = match.group(2)
+            tensors = tensors_to_vectorize.setdefault(block_key, [])
+            if len(tensors) <= block_idx:
+                tensors.extend([None] * (block_idx - len(tensors) + 1))
+            assert tensors[block_idx] is None, f"Duplicate key {k}"
+            tensors[block_idx] = v
+        else:
+            vectorized_dict[k] = v
+
+    # now we have to vectorize the tensors
+    for k, tensors in tensors_to_vectorize.items():
+        vectorized_dict[cast(str, with_prefix(prefix, k))] = jnp.stack(tensors, axis=0)
+
+    return vectorized_dict
+
+
+def _unstack_state_dict(state_dict: StateDict, prefix: Optional[str] = None) -> StateDict:
+    """
+    Unstack all keys matching prefix in a new state dict, returning a state dict that has all keys matching
+    prefix unstacked, but otherwise the same. Mostly for use with [haliax.nn.Stacked][].
+
+    Unstacked in this case means roughly "compatible with a torch.nn.Sequential", which means that the
+    keys are of the form "<prefix>.0.<key>", "<prefix>.1.<key>", etc.
+    """
+    new_dict: StateDict = {}
+    prefix = with_prefix(prefix, "")
+    assert prefix is not None
+
+    for k, v in state_dict.items():
+        if k.startswith(prefix) and is_jax_or_hax_array_like(v):
+            for i, v_i in enumerate(v):
+                new_dict[f"{prefix}{i}.{k[len(prefix):]}"] = v_i
+        else:
+            new_dict[k] = v
+
+    return new_dict

{haliax-1.4.dev336 → haliax-1.4.dev339}/src/haliax/quantization.py

@@ -14,9 +14,10 @@ import jax.random as jrandom
 from aqt.jax.v2.aqt_dot_general import DotGeneral
 from jax import numpy as jnp
 from jax.tree_util import DictKey, FlattenedIndexKey, GetAttrKey, SequenceKey
-from jax.typing import DTypeLike
+from jaxtyping import DTypeLike, PyTree
 
 import haliax.nn as hnn
+from haliax.state_dict import StateDict
 from haliax.types import PrecisionLike
 
 from ._src.fp8 import dot_general_with_precision, in_qdq, out_qdq
@@ -206,6 +207,10 @@ class Int8DotGeneralOp(OverwriteWithGradient):
         cfg = aqt_config.set_context(self.cfg, jrandom.PRNGKey(42), train_step=None)
         return cfg(lhs, rhs, dimension_numbers, precision, preferred_element_type)
 
+    def to_state_dict(tree: PyTree, prefix: Optional[str] = None) -> StateDict:
+        warnings.warn("Ignore all int8 states (if any) for now.")
+        return {}
+
 
 @dataclass(frozen=True)
 class QuantizationConfig:

{haliax-1.4.dev336 → haliax-1.4.dev339}/src/haliax/random.py

@@ -24,6 +24,7 @@ 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.dev336 → haliax-1.4.dev339}/tests/test_scan.py

@@ -1,11 +1,9 @@
 import equinox as eqx
 import jax
 import pytest
-from equinox import filter_value_and_grad
 
 import haliax as hax
-from haliax.jax_utils import tree_checkpoint_name
-from haliax.nn.scan import BlockSeq, Stacked, StackedCheckpointPolicy
+from haliax.nn.scan import BlockSeq, Stacked
 
 
 def test_unstacked():
@@ -132,6 +130,8 @@ def test_scan_with_aux_named_args():
     z_seq, z_seq_scan = m_seq.scan(x, initial_y, key=jax.random.split(jax.random.PRNGKey(2), Block.size))
     assert hax.all(hax.isclose(z, z_seq, atol=1e-5))
 
+    z_seq_scan = hax.stack(Block, z_seq_scan)
+
     assert hax.all(hax.isclose(z_scan, z_seq_scan, atol=1e-5))
 
 
@@ -164,56 +164,3 @@ def test_stacked_to_state_dict():
     y2 = m2.fold(input, key=key)
 
     assert hax.all(hax.equal(y, y2))
-
-
-def test_checkpoint_carries():
-    class Module(eqx.Module):
-        named: hax.NamedArray
-
-        def __call__(self, x):
-            y = tree_checkpoint_name(hax.sin(x + self.named), "sin")
-            y = tree_checkpoint_name(hax.cos(y + x), "cos")
-            return y + x
-
-        @staticmethod
-        def init(named):
-            return Module(named=named)
-
-    Block = hax.Axis("block", 4)
-    E = hax.Axis("E", 10)
-
-    initial_named = hax.random.uniform(jax.random.PRNGKey(0), (Block, E))
-
-    carry_policy = StackedCheckpointPolicy(save_carries=True, save_outputs=False, save_block_internals=False)
-    save_nothing = StackedCheckpointPolicy(save_carries=False, save_outputs=False, save_block_internals=False)
-    save_everything = StackedCheckpointPolicy(save_carries=True, save_outputs=True, save_block_internals=True)
-    save_internals = StackedCheckpointPolicy(save_carries=False, save_outputs=False, save_block_internals=True)
-    save_cos = StackedCheckpointPolicy(save_carries=False, save_outputs=False, save_block_internals=["cos"])
-    save_sin_carry = StackedCheckpointPolicy(save_carries=True, save_outputs=False, save_block_internals=["sin"])
-
-    for name, (policy, expected_scan_shapes) in {
-        "carry": (carry_policy, [(E.size,), (Block.size, E.size)]),
-        "nothing": (save_nothing, [(E.size,)]),
-        "everything": (save_everything, [(E.size,), (Block.size, E.size), (Block.size, E.size)]),
-        "internals": (save_internals, [(E.size,), (Block.size, E.size), (Block.size, E.size)]),
-        "cos": (save_cos, [(E.size,), (Block.size, E.size)]),
-        "sin": (save_sin_carry, [(E.size,), (Block.size, E.size), (Block.size, E.size)]),
-    }.items():
-        m = Stacked.init(
-            Block,
-            Module,
-            gradient_checkpointing=policy,
-        )(named=initial_named)
-
-        def loss_fn(m, x):
-            y = m.fold(x)
-            return hax.sum(y).scalar()
-
-        grad_fn = filter_value_and_grad(loss_fn)
-
-        jaxpr = jax.make_jaxpr(grad_fn)(m, hax.random.uniform(jax.random.PRNGKey(1), (E,)))
-        closed_call = next(eqn for eqn in jaxpr.jaxpr.eqns if eqn.primitive == jax.core.closed_call_p)
-        out_shapes = [out.aval.shape for out in closed_call.outvars]
-
-        # saved_residuals doesn't give me sensible results, so I'm doing this by hand
-        assert out_shapes == expected_scan_shapes, f"{name}: Expected {expected_scan_shapes}, got {out_shapes}"

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

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