haliax 1.4.dev331__tar.gz → 1.4.dev333__tar.gz

{haliax-1.4.dev331 → haliax-1.4.dev333}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev331
+Version: 1.4.dev333
 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.dev331 → haliax-1.4.dev333}/docs/api.md

@@ -55,11 +55,11 @@ Occasionally, an axis size can be inferred in some circumstances but not others.
 
 ### Combining Arrays
 
-We don't include `hstack` or `vstack` because we prefer semantic axes.
 
 ::: haliax.concatenate
 ::: haliax.stack
 
+(We don't include `hstack` or `vstack` because they are subsumed by `stack`.)
 
 ## Array Manipulation
 
@@ -289,6 +289,16 @@ See also the section on [Partitioning](partitioning.md).
 ::: haliax.partitioning.physical_axis_size
 ::: haliax.partitioning.sharding_for_axis
 
+
+## Gradient Checkpointing
+
+Haliax mainly just defers to JAX and [equinox.filter_checkpoint][] for gradient checkpointing. However,
+we provide a few utilities to make it easier to use.
+
+See also [haliax.nn.StackedCheckpointPolicy][].
+
+::: haliax.tree_checkpoint_name
+
 ### Old API
 
 These functions are being deprecated and will be removed in a future release.

{haliax-1.4.dev331 → haliax-1.4.dev333}/docs/broadcasting.md

@@ -62,8 +62,8 @@ This method takes one or more axes and adds them to the array.
 
 ## Explicit Broadcasting Functions
 
-::: haliax.broadcast_axis
-::: haliax.broadcast_to
+* [haliax.broadcast_axis][]
+* [haliax.broadcast_to][]
 
 
 <!--

{haliax-1.4.dev331 → haliax-1.4.dev333}/docs/css/mkdocstrings.css

@@ -1,9 +1,19 @@
 /* Indentation. */
-div.doc-contents:not(.first) {
+div.doc-contents {
   padding-left: 25px;
   border-left: .05rem solid var(--md-typeset-table-color);
 }
 
+
+div.doc-class:not(.doc-contents .doc-contents)::after {
+  content: "";
+  display: block;
+  width: 100%;
+  height: 1px; /* Adjust thickness */
+  background-color: black; /* Adjust color */
+  margin: 10px 0; /* Adjust spacing */
+}
+
 /* Mark external links as such. */
 a.external::after,
 a.autorefs-external::after {

{haliax-1.4.dev331 → haliax-1.4.dev333}/docs/fp8.md

@@ -1,20 +1,11 @@
-# FP8 Training
+# Quantized Training
 
 !!! warning
 
-    FP8 training in Haliax is currently experimental and may change in the future.
+    FP8 and Int8 training in Haliax is currently experimental and may change in the future.
 
-FP8 refers to 8-bit floating point numbers. FP8 is a massively reduced precision compared to the 32-bit floating point numbers
-or 16-bit floating point numbers that are typically used in deep learning: there are only 256 possible values in FP8, compared to
-the (almost) 2^32 in 32-bit and 2^16 in 16-bit. However, FP8 is still useful for training deep learning models, especially on
-hardware that is optimized for FP8. In particular, it can massively accelerate training on hardware that is optimized for FP8:
-H100 has 2x FP8 FLOPS compared to FP16 FLOPS and almost 60x(!) compared to F32 FLOPS.
-
-The FP8 in Haliax is currently designed to optimize throughput on FP8-enabled devices (currently H100) rather
-than to save memory. In particular, Haliax's FP8 support is not designed to quantize a model to FP8 for deployment,
-though this shouldn't be that hard to add for models that were trained using this functionality.
-We would be happy to accept contributions to add this functionality,
-and are happy to work with you to do so. In particular, adding this for models trained using Haliax's FP8 should be easy.
+Haliax supports training with FP8 and int8. This is useful for training on hardware that is optimized for FP8 or Int8,
+such as the H100 (fp8) or A100s (int8) and TPU v5 and newer (int8).
 
 ## TL;DR
 
@@ -38,15 +29,42 @@ module = haxq.apply_updates(module, updates, overwrite)
 
 And train your model like normal.
 
-## How to use FP8
+Similarly, you can use `Int8` by setting `Int8=True` in the `QuantizationConfig` object.
 
-To use FP8, you need to do two things:
 
-* Enable FP8 for the layers you want to use FP8
-* Modify your training step to be compatible with FP8
 
-Each of these is just a couple of lines of code.
+## What is FP8?
+
+FP8 refers to 8-bit floating point numbers. FP8 is a massively reduced precision compared to the 32-bit floating point numbers
+or 16-bit floating point numbers that are typically used in deep learning: there are only 256 possible values in FP8, compared to
+the (almost) 2^32 in 32-bit and 2^16 in 16-bit. However, FP8 is still useful for training deep learning models, especially on
+hardware that is optimized for FP8. In particular, it can massively accelerate training on hardware that is optimized for FP8:
+H100 has 2x FP8 FLOPS compared to FP16 FLOPS and almost 60x(!) compared to F32 FLOPS.
+
+The FP8 in Haliax is currently designed to optimize throughput on FP8-enabled devices (currently H100) rather
+than to save memory. In particular, Haliax's FP8 support is not designed to quantize a model to FP8 for deployment,
+though this shouldn't be that hard to add for models that were trained using this functionality.
+We would be happy to accept contributions to add this functionality,
+and are happy to work with you to do so. In particular, adding this for models trained using Haliax's FP8 should be easy.
+
+See this [FP8 Primer](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html) for more information on FP8.
+
+## What is Int8?
 
+Int8 refers to 8-bit integers. Int8 has the same number of bits as FP8, but the interpretation is different: instead of
+exponentially spaced numbers, Int8 has linearly spaced numbers.
+
+In Haliax, we support Int8 training through Google's [AQT](https://github.com/google/aqt) library. AQT (for
+"Accurate Quantization Training") is a library that allows you to train models with quantization-aware training (QAT).
+
+## How to use FP8 or Int8 in Haliax
+
+To use quantized training in Haliax, you need to do three things:
+
+* Enable FP8 (or int8) for the layers you want
+* Modify your training step to be compatible
+
+Each of these is just a couple of lines of code.
 
 ```python
 import haliax as hax
@@ -96,8 +114,8 @@ updates, opt_state = opt.update(grads, opt_state, params=module)  # or however y
 module = hax.quantization.apply_updates(module, updates, grads)
 ```
 
-That's it! Just a few lines of code to enable FP8. The `quantize_linear_layers` function will transform your module to use FP8
-for linear layers (or a subset if you want), and the combo of `partition_for_grad_overwrite` and `apply_updates` function will apply the updates to the module
+That's it! Just a few lines of code to enable FP8. The `quantize_linear_layers` function will transform your module to use
+quantization-aware training for linear layers (or a subset if you want), and the combo of `partition_for_grad_overwrite` and `apply_updates` function will apply the updates to the module
 in a way that is compatible with FP8.
 
 ## How FP8 works
@@ -134,6 +152,13 @@ original precision.  It remembers the maximum absolute value for each of the inp
 and scales the gradients back to the original precision. It remembers the maximum absolute value for the incoming
 gradient and stores it in the gradient.
 
+## How Int8 works
+
+Int8 is in principle the same, though the details differ. AQT is a much more flexible library than the FP8 implementation,
+because it can be a bit more finicky. We use AQT directly, and we recommend you look at the
+[AQT documentation](https://github.com/google/aqt?tab=readme-ov-file#how-aqt-works-internally) for more
+information on how it works.
+
 # API Reference
 
 ## Functions
@@ -152,7 +177,8 @@ gradient and stores it in the gradient.
 
 ::: haliax.quantization.DefaultDotGeneralOp
 ::: haliax.quantization.Fp8DotGeneralOp
+::: haliax.quantization.Int8DotGeneralOp
 
 ## Configuration
 
-::: haliax.quantization.Fp8Config
+::: haliax.quantization.QuantizationConfig

{haliax-1.4.dev331 → haliax-1.4.dev333}/docs/nn.md

@@ -5,6 +5,7 @@
 
 Haliax provides a small number of neural network modules that are compatible with Equinox, though
 they naturally all use [haliax.NamedArray][]. (We welcome PRs for more modules! Nothing too exotic though.)
+
 The most interesting of these modules is [haliax.nn.Stacked][], which allows you to create homogenous "stacks"
 of the same module (e.g. transformer blocks), which is a common pattern in deep learning.
 
@@ -21,11 +22,12 @@ of the same module (e.g. transformer blocks), which is a common pattern in deep
 
 ### Meta
 
-::: haliax.nn.scan.BlockFoldable
-::: haliax.nn.BlockSeq
-::: haliax.nn.Stacked
 ::: haliax.nn.MLP
 
+### Stacked
+
+See the full documentation of [Stacked](stacked.md).
+
 ### Convolution
 
 Unlike other frameworks, Haliax doesn't distinguish between 1D, 2D, and 3D, and general convolutions. Instead, we have

haliax-1.4.dev333/docs/stacked.md

@@ -0,0 +1,203 @@
+# Module Stacks
+
+A core pattern for larger models in JAX is the "scan-over-layers" pattern, where you have a sequence of layers
+that get stacked together, and you use [jax.lax.scan][] or [haliax.fold][] or [haliax.scan][] to apply them to a
+sequence of inputs. In Haliax, layers are represented as [equinox.nn.Module][]s, and the [haliax.nn.Stacked][] module
+provides a way to create a sequence of layers that can be applied to a sequence of inputs that implements the
+scan-over-layers pattern.
+
+## Stacked
+
+[haliax.nn.Stacked][] lets you apply a layer sequentially to an input, scanning over a "Layers" axis. For instance,
+a Transformer might use a Stacked for its Transformer blocks:
+
+
+```python
+class TransformerBlock(eqx.Module):
+
+    def __init__(self, config: TransformerConfig, layer_index, *, key):
+        attn_key, mlp_key = jax.random.split(key)
+        self.attention = Attention.init(config, key=attn_key)
+        self.mlp = MLP.init(config, key=mlp_key)
+        self.ln1 = LayerNorm.init(config.Hidden)
+        self.ln2 = LayerNorm.init(config.Hidden)
+        self.layer_index = layer_index
+
+    def __call__(self, x):
+        y = self.attention(self.ln1(x))
+        x = x + y
+        y = self.mlp(self.ln2(x))
+        return x + y
+
+class Transformer(eqx.Module):
+    def __init__(self, config: TransformerConfig):
+        self.blocks = Stacked.init(Layers, TransformerBlock)(
+            config,  # static configuration
+            scale=hax.arange(Layers),  # dynamic configuration. Each layer gets a scalar scale value [0, 1, 2, ...]
+            key=jax.random.split(key, Layers),  # dynamic configuration. Each layer gets a different key
+        )
+
+    def __call__(self, x: NamedArray) -> NamedArray:
+        # morally the equivalent of:
+        # for block in self.blocks:
+        #     x = block(x)
+        # Except that it works better with JAX compile times.
+
+        return self.blocks.fold(x)
+```
+
+You can think of [haliax.nn.Stacked][] as an analog to PyTorch's
+[torch.nn.Sequential](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html), except that
+every layer in the sequence must have exactly the same shape and configuration.
+
+Internally, a Stacked is a single copy of the module, except that every NamedArray inside that module
+has a Block axis prepended (as though they were stacked with [haliax.stack][]). Similarly, every JAX array
+inside the module has its first axis prepended with an axis of the same size as the Block axis, as though
+they were stacked with [jax.numpy.stack][].
+
+When you call the Stacked, it scans over the Block axis, applying the module to each element of the Block.
+
+### Creating a Stacked
+
+To create a Stacked, we provide `Stacked.init`, which takes a "Layers" [haliax.Axis][] and another Module as
+well as args and kwargs for that module. The Layer is the axis that the Stacked will scan over, and the `args`
+and `kwargs` are implicitly vmapped over the Layers.
+
+For instance, to create a stack of GPT2 blocks, you might do:
+
+```python
+import jax.random
+
+blocks = Stacked.init(Layers, Gpt2Block)(
+    config,  # static configuration
+    scale=hax.arange(Layers),  # dynamic configuration. Each layer gets a scalar scale value [0, 1, 2, ...]
+    key=jax.random.split(key, Layers.size),  # dynamic configuration. Each layer gets a different key
+)
+```
+
+Any NamedArray passed to the Stacked init will have its Layers axis (if present) vmapped over. Any
+JAX array will have its first axis vmapped over.
+
+
+### Fold Blocks vs Scan Blocks
+
+The Stacked module provides two ways to apply the layers: `fold` and `scan`.  A fold is the moral equivalent of this for loop:
+
+```python
+for block in self.blocks:
+    x = block(x)
+```
+
+while a scan is the moral equivalent of this for loop:
+
+```python
+out = []
+for block in self.blocks:
+    x, y = block(x)
+    out.append(y)
+
+return x, stack(out)
+```
+
+Blocks can be coded to either support fold or scan, but not both.
+A "fold Block" should have the signature `def __call__(self, x: Carry) -> Carry`,
+while a "scan Block" should have the signature `def __call__(self, x: Carry) -> Tuple[Carry, Output]`.
+
+(See also [jax.lax.scan][], [haliax.fold][], and [haliax.scan][].)
+
+### Requirements for Stacked Blocks
+
+As we said above, the Stacked module requires that all the layers have the same shape and configuration.
+
+A further constraint is that the elements of the stack must have the same Python control flow. This is the usual
+constraint imposed on jit-compiled functions in JAX. All control flow must use `jax.lax` primitives like
+[jax.lax.cond][], [jax.lax.while_loop][], and [jax.lax.scan][]. You can't use Python control flow like `if` or `for`
+except for static control flow that is the same for all elements of the stack.
+
+## BlockSeq and BlockFoldable
+
+We also provide a way to create a sequence of layers that can be applied to a sequence of inputs that implements the
+same interface as [haliax.nn.Stacked][], but with a different implementation. This is the [haliax.nn.BlockSeq][] module.
+BlockSeq implements those for loops directly, rather than using [haliax.fold][] or [haliax.scan][].
+
+[haliax.nn.scan.BlockFoldable][] is an interface that both [haliax.nn.Stacked][] and [haliax.nn.BlockSeq][] implement.
+
+## Gradient Checkpointing
+
+The [haliax.nn.Stacked][] module also provides a way to do gradient checkpointing, which can be useful for deep models.
+
+Gradient checkpointing, aka rematerialization, is a technique for trading off memory usage for compute time.
+Instead of storing all the intermediate activations of a model, you store only a subset and recompute the rest
+as needed. (XLA automatically recomputes the rest for you as needed.)
+
+[JAX's checkpointing mechanism]((https://docs.jax.dev/en/latest/gradient-checkpointing.html) is highly flexible,
+and we provide a relatively simple interface to it for use with `Stacked`.
+
+### Simple Checkpointing
+In the simplest case, you can enable a usually-good-enough checkpointing policy by passing `gradient_checkpointing=True`
+to the `Stacked.init` call:
+
+```python
+blocks = Stacked.init(Layers, TransformerBlock, gradient_checkpointing=True)(
+    config,
+    scale=hax.arange(Layers),
+    key=jax.random.split(key, Layers.size),
+)
+```
+
+This will preserve the intermediate "carries" and the "outputs" of the scans, while rematerializing (i.e. recomputing)
+the rest of the computation as needed during backpropagation.
+
+### Custom Checkpointing Policies
+
+If you need more control over the checkpointing policy, you can pass a [haliax.nn.StackedCheckpointPolicy][] object to
+the Stacked init:
+
+```python
+policy = StackedCheckpointPolicy(
+   save_carries=True,  # default
+   save_outputs=True,  # default
+   save_intermediates=False,  # default
+)
+```
+
+### Saving Block-Internal Values
+
+"`intermediates`" refers to the internal computation of the block. If you set `save_intermediates=True`, then
+all internals of every block will be saved. This can be expensive.
+
+You can also pass a list of strings to `save_intermediates` to specify which intermediates to save.
+
+You could, for instance, save the output of the attention layer using [haliax.tree_checkpoint_name][]:
+
+```python
+class TransformerBlock(eqx.Module):
+    def __call__(self, x):
+        y = self.attention(self.ln1(x))
+        y = haliax.tree_checkpoint_name(y, "attn_out")
+        x = x + y
+        y = self.mlp(self.ln2(x))
+        return x + y
+
+policy = StackedCheckpointPolicy(save_carries=True, save_block_internals=["attn_out"])
+```
+
+With this policy, the output of the attention layer will be saved during the forward pass.
+
+This will save an extra attention computation in the backward pass, adding $`O(N * Pos * Hidden)`$ memory usage,
+which is double that required by the default policy.
+
+
+### Offloading Checkpointed Values
+
+Both `save_carries` and `save_outputs` can either be a boolean or the string "offload". If "offload", then the
+checkpointed values will be offloaded to the host during the forward pass, and reloaded during the backward pass.
+
+
+## API
+
+::: haliax.nn.Stacked
+::: haliax.nn.BlockSeq
+::: haliax.nn.scan.BlockFoldable
+
+::: haliax.nn.StackedCheckpointPolicy

{haliax-1.4.dev331 → haliax-1.4.dev333}/docs/state-dict.md

@@ -157,4 +157,21 @@ to join the prefix to the keys of the state dict.
 
 ## API Reference
 
-::: haliax.state_dict
+::: haliax.state_dict.StateDict
+::: haliax.state_dict.ModuleWithStateDictSerialization
+
+### Saving and Loading State Dicts
+::: haliax.state_dict.save_state_dict
+::: haliax.state_dict.load_state_dict
+
+### Converting betweewn State Dicts and Modules
+
+::: haliax.state_dict.from_state_dict
+::: haliax.state_dict.to_state_dict
+
+### Torch Compatibility
+
+::: haliax.state_dict.from_torch_compatible_state_dict
+::: haliax.state_dict.to_torch_compatible_state_dict
+::: haliax.state_dict.flatten_linear_layers
+::: haliax.state_dict.unflatten_linear_layers

{haliax-1.4.dev331 → haliax-1.4.dev333}/mkdocs.yml

@@ -43,25 +43,34 @@ plugins:
             - https://jax.readthedocs.io/en/latest/objects.inv
             - https://docs.kidger.site/equinox/objects.inv
             - https://einops.rocks/objects.inv
+          selection:
+            inherited_members: true
           options:
             docstring_options:
               ignore_init_summary: true
               docstring_style: sphinx
             show_source: false
-            docstring_section_style: list
-            heading_level: 5
+            filters:
+              - "!^_"
+            heading_level: 4
             inherited_members: true
+            members_order: source
             merge_init_into_class: true
+            parameter_headings: true
+            separate_signature: false
             load_external_modules: true
             preload_modules: [haliax, haliax.core]
-#            separate_signature: true
+            show_if_no_docstring: true
             show_root_heading: true
             show_root_full_path: false
-#            show_signature_annotations: true
-            show_symbol_type_heading: false
+            show_signature_annotations: true
+            docstring_section_style: list
+            show_symbol_type_heading: true
             show_symbol_type_toc: false
             signature_crossrefs: true
             line_length: 100
+            summary: true
+
   - include-markdown
 extra_css:
   - css/material.css
@@ -85,6 +94,7 @@ 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.dev333/src/haliax/__about__.py

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

{haliax-1.4.dev331 → haliax-1.4.dev333}/src/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-1.4.dev331 → haliax-1.4.dev333}/src/haliax/_src/state_dict.py

@@ -395,7 +395,7 @@ def flatten_linear_layers(tree: T) -> T:
         new_In: Axis = flatten_axes(layer.In, "__IN__")
 
         if weight.array is not None:
-            out_first = layer.out_first
+            out_first = layer._out_first
             weight = weight.flatten_axes(layer.Out, new_Out).flatten_axes(layer.In, new_In)
 
             if out_first:

{haliax-1.4.dev331 → haliax-1.4.dev333}/src/haliax/axis.py

@@ -35,6 +35,8 @@ def make_axes(**kwargs: int) -> Tuple[Axis, ...]:
     Example:
     ```
     X, Y = axes(X=10, Y=20)
+    ```
+
     """
     return tuple(Axis(name, size) for name, size in kwargs.items())
 

{haliax-1.4.dev331 → haliax-1.4.dev333}/src/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-1.4.dev331 → haliax-1.4.dev333}/src/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: