Scan and Fold¤

A common pattern in deep learning is to apply a sequence of layers to an input, feeding the output from one layer to the next. In JAX, this is often done with jax.lax.scan.

As the docs say, scan does an operation sort of like this in Python:

def scan(f, init, xs, length=None):
  if xs is None:
    xs = [None] * length
  carry = init
  ys = []
  for x in xs:
    carry, y = f(carry, x)
    ys.append(y)
  return carry, np.stack(ys)

Haliax provides two versions of this pattern: haliax.fold and haliax.scan. haliax.scan works much like JAX's scan, except it is curried and it works with NamedArrays. haliax.fold is a more restricted version of scan that is easier to use if you don't need the full generality of scan. (It works with functions that only return carry, not carry, output.)

haliax.scan ¤

Unlike JAX's scan, Haliax's scan is curried - it takes the function and configuration first, then the initial carry and scan arguments as a separate call: scan(f, axis)(init, xs).

Key Features¤

Works with named axes using haliax.NamedArray
Supports gradient checkpointing for memory efficiency, including several advanced checkpointing policies
Integrates with equinox.Module for building neural networks

Basic Example¤

Here's a practical example of using haliax.scan to sum values along an axis while keeping track of intermediates:

Time = Axis("Time", 100)
Features = Axis("Features", 16)

# Create time series data
data = hax.random.normal(PRNGKey(0), (Time, Features))

def running_stats(state, x):
    count, mean, min_val, max_val = state
    count += 1
    # this is a common pattern to improve the robustness of the mean calculation
    delta = x - mean
    mean = mean + delta / count
    min_val = hax.minimum(min_val, x)
    max_val = hax.maximum(max_val, x)

    return (count, mean, min_val, max_val), mean


# Initialize state: (count, mean, min, max)
init_state = (
    0.0,
    hax.zeros((Features,)),
    hax.full((Features,), float('inf')),
    hax.full((Features,), float('-inf'))
)

final_state, running_means = hax.scan(running_stats, Time)(init_state, data)

Note that:

scan is curried: scan(f, axis)(init, xs)
running_stats returns a tuple of (carry, output), which is why we have two return values from scan
the running_means will have shape (Time, Features), with the mean at each time step
the final_state will have the same shape as the initial state

Using `scan` with no inputs¤

You can also use scan without any inputs if you want:

Time = Axis("Time", 100)
Features = Axis("Features", 16)

def simulate_brownian_motion(state, _):
    return state + hax.random.normal(PRNGKey(0), Features), state

init_state = hax.zeros((Features,))

final_state, path = hax.scan(simulate_brownian_motion, Time)(init_state, None)

More commonly, you might use this for an RNN or Transformer model. (See haliax.nn.Stacked.)

haliax.fold ¤

haliax.fold is a simpler version of haliax.scan that is easier to use when you don't need the full generality of scan. Specifically, fold is for functions that only return a carry, not a carry, output.

Morally, fold is like this Python code:

def fold(f, init, xs):
  carry = init
  for x in xs:
    carry = f(carry, x)
  return carry

Basic Example¤

Same example, but we only care about the final state:

Time = Axis("Time", 100)
Features = Axis("Features", 16)

# Create time series data
data = hax.random.normal(PRNGKey(0), (Time, Features))

def running_stats(state, x):
    count, mean, min_val, max_val = state
    count += 1
    # this is a common pattern to improve the robustness of the mean calculation
    delta = x - mean
    mean = mean + delta / count
    min_val = hax.minimum(min_val, x)
    max_val = hax.maximum(max_val, x)

    return (count, mean, min_val, max_val)


# Initialize state: (count, mean, min, max)
init_state = (
    0.0,
    hax.zeros((Features,)),
    hax.full((Features,), float('inf')),
    hax.full((Features,), float('-inf'))
)

final_state = hax.fold(running_stats, Time)(init_state, data)

haliax.map ¤

haliax.map is a convenience function that applies a function to each element of an axis. It is similar to jax.lax.map but works with NamedArrays, providing a similar interface to haliax.scan and haliax.fold.

Time = Axis("Time", 100)

data = hax.random.normal(PRNGKey(0), (Time,))

def my_fn(x):
    return x + 1

result = hax.map(my_fn, Time)(data)

You should generally prefer to use haliax.vmap instead of haliax.map, but it's there if you need it. (It uses less memory than haliax.vmap but is slower.)

Gradient Checkpointing / Rematerialization¤

Both haliax.scan and haliax.fold support gradient checkpointing, which can be useful for deep models. Typically, you'd use this as part of haliax.nn.Stacked or haliax.nn.BlockSeq but you can also use it directly.

Gradient checkpointing is a technique for reducing memory usage during backpropagation by recomputing some intermediate values during the backward pass. This can be useful when you have a deep model with many layers.

TL;DR Guidance¤

Here is some guidance on when to use gradient checkpointing:

Use remat=False if you need to reduce computation and have lots of memory. This is the default in haliax.scan.
Use remat=True for most models. It's usually good enough. This is the default in haliax.nn.Stacked.
Use remat="nested" if you need to reduce memory usage.
Use save_block_internals sparingly, but it is your best tool for trading increased memory usage for reduced computation if you need something between remat=True and remat=False.
Use save_carries="offload" if you need to reduce memory usage at the cost of recomputation. This is a new feature in JAX and doesn't seem to reliably work yet.

Simple Checkpointing¤

In the simplest case, you can enable a usually-good-enough checkpointing policy by passing remat=True:

final_state = hax.fold(running_stats, Time, remat=True)(init_state, data)

("remat" is short for "rematerialization", which is another term for gradient checkpointing.)

This will preserve the intermediate "carries" and other inputs the fold function needs, while rematerializing (i.e. recomputing) the internal state of each block (i.e. call to the running_stats function) as needed during backpropagation.

Nested Scan¤

Simple checkpointing requires O(N) memory where \(N\) is the number of blocks. A nested scan lets you reduce this to O(sqrt(N)) memory, at the cost of a bit more computation. You can enable this by passing remat="nested":

final_state = hax.fold(running_stats, Time, remat="nested")(init_state, data)

This will break the scan into a double loop, where the outer loop has sqrt(N) blocks and the inner loop has sqrt(N) blocks (with appropriate rounding).

Functionally, it does something like:

outer_size = int(sqrt(N))  # ensuring outer_size divides N
blocks = haliax.rearrange("block -> (outer inner)", blocks, outer=outer_size)

state = init_state
for o in range(outer_size):
    inner_blocks = blocks["outer", o]

    for i in range(inner_size):
        state = f(state, inner_blocks["inner", i])

    # not real jax
    state = save_for_backward(state)

where we save only the carries from the outer loop, and fully rematerialize the inner loop.

If C is the amount of memory needed for the carry, and N is the number of blocks, then the memory usage of the nested scan is 2 * C * sqrt(N). In addition, you need enough memory to do backward in one block.

In practice, nested scan is about 20% slower than simple checkpointing (for Transformers), but uses much less memory.

Advanced: customizing the number of blocks¤

You can also customize the number of blocks in the outer loop by using a policy:

policy = ScanCheckpointPolicy(nested=4)  # 4 outer blocks

Note that by itself this doesn't help you at all except potentially requiring more memory. You can potentially combine it with other policy options to make things faster though.

Custom Checkpointing Policies¤

If you need more control over the checkpointing policy, you can pass a haliax.nn.ScanCheckpointPolicy object to the scan or fold call:

policy = ScanCheckpointPolicy(
   save_carries=True,  # default
   save_inputs=True,  # default
   save_block_internals=False,  # default
)

Saving Block-Internal Values¤

"internals" refers to the internal computation of the block. If you set save_block_internals=True, then all internals of every block will be saved. This can be expensive and mostly negates the benefits of checkpointing.

Instead you can choose which internals to save by passing a list of strings to save_block_internals:

def my_complex_fn(state, x):
    y = x + state
    y = hax.sin(y) + x
    y = hax.tree_checkpoint_name(y, "y")
    y = hax.cos(y) + x
    y = hax.tree_checkpoint_name(y, "z")
    return y

policy = ScanCheckpointPolicy(save_carries=True, save_block_internals=["y"])

final_state = hax.fold(my_complex_fn, Time, remat=policy)(init_state, data)

With this policy, the output of the sin function will be saved during the forward pass.

This will save an extra sin computation in the backward pass, adding \(O(N * Pos * Hidden)\) memory usage, which is double that required by the default policy, but it reduces the amount of recomputation needed. (It's probably not worth it in this case.)

Offloading Checkpointed Values¤

Both save_carries and save_inputs 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.

In addition, you can offload block internals by passing a list of strings to offload_block_internals:

policy = ScanCheckpointPolicy(save_carries=True, save_block_internals=["y"], offload_block_internals=["z"])

Summary of String and Boolean Aliases¤

remat=True is the same as remat=ScanCheckpointPolicy(save_carries=True, save_inputs=True)
remat="full" is the same as remat=True
remat=False is the same as remat=ScanCheckpointPolicy(disable=True)
remat="nested" is the same as remat=ScanCheckpointPolicy(nested=True)
remat="offload" is the same as remat=ScanCheckpointPolicy(save_carries="offload", save_inputs="offload")
remat="save_all" is the same as remat=ScanCheckpointPolicy(save_carries=True, save_inputs=True, save_block_internals=True), which should be the same as not using remat at all...

Memory and Computation Tradeoffs¤

Let N be the number of blocks, C be the memory needed for the carry, and I be the internal memory needed for each block. Let F be the amount of computation needed for each block. Constants are added for a bit more precision but are not exact. This is assuming that backward requires ~twice the flops as forward, which is roughly right for Transformers.

Policy	Memory Usage	Computation
`remat=False`	`O(N * C + N * I)`	`O(3 * N * F)`
`remat=True`	`O(N * C + I)`	`O(4 * N * F)`
`remat="nested"`	`O(2 * sqrt(N) * C + I)`	`O(5 * N * F)`

(Which shows why nested scan is about 20% slower than simple checkpointing. The math says 25% but it's more like 20% in practice.) Any nested remat will require 5 * N * F computation, which is about 25% more than simple remat.

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:

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, 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:

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.

Apply Blocks in Parallel with `vmap`¤

Sometimes you may want to apply each block independently, without feeding the output of one block into the next. Stacked.vmap does exactly that: it uses haliax.vmap to broadcast the initial value to every block and evaluates them in parallel, returning the stack of outputs.

y = stacked.vmap(x)

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:

for block in self.blocks:
    x = block(x)

while a scan is the moral equivalent of this for loop:

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. It exposes the usual fold and scan methods as well as helpers fold_via and scan_via which return callables that perform the respective operations using a custom block function.

API¤

`fold(fn: Callable, axis: AxisSelector, *, remat: ScanCheckpointSpec = False, reverse: bool = False, unroll: int = 1, is_scanned: BoolAxisSpec = is_named_or_shaped_array_like) -> Callable` ¤

fold(fn: Callable[[Carry, X], Carry], axis: AxisSelector, *, remat: ScanCheckpointSpec = False, reverse: bool = False, unroll: int = 1, is_scanned: BoolAxisSpec = is_named_or_shaped_array_like) -> Callable[[Carry, PyTree[X]], Carry]

fold(fn: Callable, axis: AxisSelector, *, remat: ScanCheckpointSpec = False, reverse: bool = False, unroll: int = 1, is_scanned: BoolAxisSpec = is_named_or_shaped_array_like) -> Callable

Slightly simpler implementation of scan that folds over the named axis of the array, not returning intermediates.

As with scan, this function is curried: it takes the function, axis, and configuration arguments first, and then the initial carry and then any arguments to scan over as a separate curried function call.

Unnamed arrays will have their first axis scanned over, unless they are scalars, in which case they will be passed through unchanged.

Parameters:

fn ¤
(Callable) –

function to reduce over
axis ¤
(AxisSelector) –

axis to reduce over
reverse ¤
(bool, default: False ) –

if True, reduce in reverse
unroll ¤
(int, default: 1 ) –

unroll the loop by this amount
nested_scan ¤
–

Use nested scans to reduce memory usage. If an integer, use that many nested scans. If true, use the closest int to sqrt(axis.size) that divides axis.size, which gives O(sqrt(N)) memory and O(N) time.
is_scanned ¤
(BoolAxisSpec, default: is_named_or_shaped_array_like ) –

a function that takes a leaf of the tree and returns True if it should be scanned over, False otherwise. Behaves similarly to the default argument in filter_jit

Returns:

Callable –

A function that takes the initial carry and then the arguments to reduce over, and returns the final carry

`scan(f: Callable, axis: AxisSelector, *, remat: ScanCheckpointSpec = False, reverse=False, unroll=1, is_scanned: BoolAxisSpec = is_named_or_shaped_array_like)` ¤

scan(f: Callable[[Carry, X], tuple[Carry, Y]], axis: AxisSelector, *, remat: ScanCheckpointSpec = False, reverse: bool = False, unroll: int = 1, is_scanned: BoolAxisSpec = is_named_or_shaped_array_like) -> Callable[[Carry, PyTree[X]], tuple[Carry, PyTree[Y]]]

scan(f: Callable, axis: AxisSelector, *, remat: ScanCheckpointSpec = False, reverse: bool = False, unroll: int = 1, is_scanned: BoolAxisSpec = is_named_or_shaped_array_like) -> Callable

Scan over a named axis. Non-scalar unnamed arrays will have their first axis scanned over.

Unlike jax.lax.scan, this function is curried: it takes the function, axis, and configuration arguments first, and then the initial carry and then any arguments to scan over as a separate curried function call.

That is, scan(f, axis)(init, xs) is equivalent to jax.lax.scan(f, init, xs)

Parameters:

f ¤
(Callable) –

function to scan over
axis ¤
(AxisSelector) –

axis to scan over
reverse ¤
(bool, default: False ) –

if True, scan in reverse
unroll ¤
(int, default: 1 ) –

unroll the loop by this amount
remat ¤
(ScanCheckpointSpec, default: False ) –

a ScanCheckpointPolicy or a boolean. If True, rematerialize block internals during gradient computation. If False, rematerialize nothing. If "offload", offload all carries and inputs. See haliax.ScanCheckpointPolicy for more information.
is_scanned ¤
(BoolAxisSpec, default: is_named_or_shaped_array_like ) –

a function that takes a leaf of the tree and returns True if it should be scanned over, False otherwise. Behaves similarly to the default argument in filter_jit

`map(fn: Callable[[X], Y], axis: Axis, *, remat: ScanCheckpointSpec = False, reverse: bool = False, unroll: int = 1, is_mapped: BoolAxisSpec = is_jax_or_hax_array_like) -> Callable[[PyTree[X]], PyTree[Y]]` ¤

NamedArray aware version of jax.lax.map. Normal arrays are mapped according to the specs as in equinox.filter_map, except that the output axis is always 0 b/c it's annoying to make anything else work.

You'll typically want to use map (instead of a vmap or just vectorized code) when you want to encourage XLA to loop over the axis to control memory.

Parameters:

fn ¤
(Callable) –

function to map over
axis ¤
(Axis) –

axis to map over
reverse ¤
(bool, default: False ) –

if True, map in reverse
unroll ¤
(int, default: 1 ) –

unroll the loop by this amount

`ScanCheckpointPolicy(save_carries: bool | Literal['offload'] = True, save_inputs: bool | Literal['offload'] = False, save_block_internals: bool | list[str] = False, offload_block_internals: list[str] = list(), prevent_cse: bool = False, disable: bool = False, simple: bool = False, nested: bool | int = False)` `dataclass` ¤

A class that represents a gradient checkpoint policy for blocks in a Stacked module. This is used to control gradient checkpointing in haliax.scan, haliax.fold, 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.

Without checkpointing, we will save all block inputs and outputs, and use them to compute the gradient. This requires memory, of course. With checkpointing, we can save only some of the computation. Typically, to compute the gradient of the scan, we need the inputs to each block, the intermediates within each block, as well as the carries (which are inputs to the next block). We can save all of these, or we can save only some of them.

With this class, you can specify if you want to save the carries, or the internals/inputs. You can also offload the carries to the host, which can reduce memory usage on the device.

Nested Remat¤

Alternatively, we can do something a bit more clever. We can break the computation into "blocks" of size "B", and save the carries and outputs of each block. Then, during the backward pass, we can recompute the outputs of each block using the carries and inputs, and then compute the gradient as usual. This requires O(B) memory and O(N) time. When B = sqrt(N), this is O(sqrt(N)) memory and O(N) time. This is the "nested scan" policy. In practice, this is about 20% slower than the O(N) memory policy, but it can be worth it for large models.

Offloading¤

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.

`save_carries: bool | Literal['offload'] = True` `class-attribute` `instance-attribute` ¤

Whether to save all carries in the forward pass. If True, (input) 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

If False, carries are recomputed in the backward pass.

`save_inputs: bool | Literal['offload'] = False` `class-attribute` `instance-attribute` ¤

Whether to save all non-carry inputs in the forward pass. If True, inputs are saved in the forward pass and used in the backward pass. If "offload", inputs are saved in the forward pass and offloaded to the host.

`save_block_internals: bool | list[str] = False` `class-attribute` `instance-attribute` ¤

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.

`offload_block_internals: list[str] = dataclasses.field(default_factory=list)` `class-attribute` `instance-attribute` ¤

List of named block internals to offload to the host. This is useful for reducing memory usage on the device while still avoiding rematerialization.

`prevent_cse: bool = False` `class-attribute` `instance-attribute` ¤

Whether to prevent common subexpression elimination in the checkpointed function.

`disable: bool = False` `class-attribute` `instance-attribute` ¤

Whether to disable gradient checkpointing entirely. This is useful for debugging.

`simple: bool = False` `class-attribute` `instance-attribute` ¤

Whether to use the simple gradient checkpointing policy. This is useful for debugging.

`nested: bool | int = False` `class-attribute` `instance-attribute` ¤

Allows for nested remat with a double scan. We reshape the stack into [nested_remat, -1] and then scan over both in sequence. If True, we find the closest int to sqrt(len(stack)) such that len(stack) % int == 0. If False, we don't do anything.

`from_bool_or_str(remat_policy: bool | str)` `staticmethod` ¤

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
"nested": use nested remat. Equivalent to True

`checkpoint(carry_name: str, input_name: str, callable)` ¤

Modules¤

`Stacked` ¤

Bases: 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 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. If your blocks are independent of each other you can instead use :py:meth:Stacked.vmap to apply every block in parallel.

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 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".

Example

>>> 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))

Methods:

flatten_for_export –

Flatten articulated named arrays for export to torch. In general this method should, for a linear layer, flatten
unflatten_from_export –

Unflatten the module after import from torch.
init –

Initialize a Stacked module. This method is curried: you can pass in the Block and module, and it will return
scan –

Scan over the stacked module. This is the same as a for loop that applies each instance of the module in sequence
fold –

Fold over the stacked module. This is the same as a for loop that applies each instance of the module in sequence
fold_via –

Return a function that folds over the stack using fn.
scan_via –

Return a function that scans over the stack using fn.
vmap_via –

Return a function that applies each block independently using fn.
vmap –

Apply each block independently using :func:haliax.vmap.
unstacked –

Returns the unstacked version of this module. This is useful for logging or saving checkpoints.
get_layer –

Return the indexth layer of this stacked module.
to_state_dict –
from_state_dict –

Attributes:

stacked (M) –
Block (Axis) –
gradient_checkpointing (ScanCheckpointPolicy) –
Layers (Axis) –

Alias for :attr:Block used by some downstream code.

`stacked: M` `instance-attribute` ¤

`Block: Axis = eqx.field(static=True)` `class-attribute` `instance-attribute` ¤

`gradient_checkpointing: ScanCheckpointPolicy = eqx.field(static=True)` `class-attribute` `instance-attribute` ¤

`Layers: Axis` `property` ¤

Alias for :attr:Block used by some downstream code.

`flatten_for_export() -> Mod` ¤

Flatten articulated named arrays for export to torch. In general this method should, for a linear layer, flatten all input axes into a single axis, and all output axes into a single axis. You can do whatever else you want to support pytorch-compatible serialization if you want.

This method is less general than to_state_dict and is only called when using to_torch_compatible_state_dict.

`unflatten_from_export(template: Mod) -> Mod` ¤

Unflatten the module after import from torch.

Template has the proper structure (e.g. articulated named axes) but the values are meaningless.

`init(Block: Axis, module: Type[M], *, gradient_checkpointing: ScanCheckpointSpec = False, prevent_cse: bool | None = None) -> ModuleInit[Stacked[M]]` `classmethod` ¤

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.

Parameters:

Block ¤
(Axis) –

The axis that will be stacked over. This is typically a "layer" axis, but could be any axis.
module ¤
(Type[M]) –

The module that will be stacked. This module must take a batched input and return a batched output.
gradient_checkpointing ¤
(ScanCheckpointSpec, default: False ) –

Whether to use gradient checkpointing. If True, uses the default policy. If a string, uses the policy specified by the string. If a ScanCheckpointPolicy, uses that policy.
prevent_cse ¤
(bool | None, default: None ) –

Whether to prevent common subexpression elimination in the checkpointed function. This is useful for debugging, but may slow down the function.

`scan(init, *extra_args, unroll: int | bool | None = None, **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 outputs as well as the final output.

That is, it behaves similarly to the following Python code:

carry = init
outputs = []

for block in self.stacked:
    carry, extra = block(carry)
    outputs.append(extra)

return carry, hax.stack(Block, outputs)

Parameters:

init ¤
–
*extra_args ¤
–
**extra_kwargs ¤
–

Returns:

`fold(init, *args, unroll: int | bool | None = None, **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:

carry = init
for block in self.stacked:
    carry = block(carry)

return carry

Parameters:

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:

`fold_via(fn: Callable[..., CarryT], *, unroll: int | bool | None = None)` ¤

fold_via(fn: FoldFunction[M, P, CarryT], *, unroll: int | bool | None = None) -> Callable[Concatenate[CarryT, P], CarryT]

fold_via(fn: Callable[[M, CarryT], CarryT], *, unroll: int | bool | None = None) -> Callable[[CarryT], CarryT]

Return a function that folds over the stack using fn.

fn should take a block and a carry and return a new carry. The returned function mirrors :func:haliax.fold over the block axis.

`scan_via(fn: Callable[..., tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None)` ¤

scan_via(fn: ScanFunction[M, CarryT, P, OutputT_co], *, unroll: int | bool | None = None) -> Callable[Concatenate[CarryT, P], tuple[CarryT, OutputT_co]]

scan_via(fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None) -> Callable[[CarryT], tuple[CarryT, OutputT_co]]

Return a function that scans over the stack using fn.

fn should take a block and a carry and return (carry, output). Semantics match :func:haliax.scan over the block axis.

`vmap_via(fn: Callable[..., OutputT_co]) -> Callable[..., OutputT_co]` ¤

vmap_via(fn: VmapFunction[M, P, OutputT_co]) -> Callable[P, OutputT_co]

vmap_via(fn: Callable[[M], OutputT_co]) -> Callable[[], OutputT_co]

Return a function that applies each block independently using fn.

fn should take a block and a carry and return an output. The returned function mirrors :func:haliax.vmap over the block axis.

`vmap(*extra_args, **extra_kwargs)` ¤

Apply each block independently using :func:haliax.vmap.

Returns the stacked outputs of each block.

`unstacked() -> 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

`get_layer(index: int) -> M` ¤

Return the indexth layer of this stacked module.

`to_state_dict(prefix: str | None = None) -> StateDict` ¤

`from_state_dict(state_dict: StateDict, prefix: str | None = None) -> M` ¤

`BlockSeq` ¤

Bases: 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).

Methods:

flatten_for_export –

Flatten articulated named arrays for export to torch. In general this method should, for a linear layer, flatten
unflatten_from_export –

Unflatten the module after import from torch.
init –

This is a curried init method that takes the Block and module and returns a function that takes
scan –
fold –
fold_via –

Return a function that folds over the sequence using fn.
scan_via –

Return a function that scans over the sequence using fn.
vmap_via –

Return a function that applies each block independently using fn.
unstacked –
from_state_dict –
to_state_dict –

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.".
get_layer –

Return the indexth block in this sequential container.

Attributes:

blocks (Sequence[M]) –
Block (Axis) –
gradient_checkpointing (ScanCheckpointPolicy) –

`blocks: Sequence[M]` `instance-attribute` ¤

`Block: Axis = eqx.field(static=True)` `class-attribute` `instance-attribute` ¤

`gradient_checkpointing: ScanCheckpointPolicy = eqx.field(static=True)` `class-attribute` `instance-attribute` ¤

`flatten_for_export() -> Mod` ¤

Flatten articulated named arrays for export to torch. In general this method should, for a linear layer, flatten all input axes into a single axis, and all output axes into a single axis. You can do whatever else you want to support pytorch-compatible serialization if you want.

This method is less general than to_state_dict and is only called when using to_torch_compatible_state_dict.

`unflatten_from_export(template: Mod) -> Mod` ¤

Unflatten the module after import from torch.

Template has the proper structure (e.g. articulated named axes) but the values are meaningless.

`init(Block: Axis, module: Type[M], *, gradient_checkpointing: ScanCheckpointSpec = False, prevent_cse: bool | None = None) -> ModuleInit[S]` `classmethod` ¤

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.

`scan(init: T, *extra_args, unroll: int | bool | None = None, **extra_kwargs)` ¤

`fold(init: T, *args, unroll: int | bool | None = None, **kwargs) -> T` ¤

`fold_via(fn: Callable[..., CarryT], *, unroll: int | bool | None = None)` ¤

fold_via(fn: FoldFunction[M, P, CarryT], *, unroll: int | bool | None = None) -> Callable[Concatenate[CarryT, P], CarryT]

fold_via(fn: Callable[[M, CarryT], CarryT], *, unroll: int | bool | None = None) -> Callable[[CarryT], CarryT]

Return a function that folds over the sequence using fn.

fn should take a block and a carry and return a new carry. The returned function mirrors :func:haliax.fold over the block axis.

`scan_via(fn: Callable[..., tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None)` ¤

scan_via(fn: ScanFunction[M, CarryT, P, OutputT_co], *, unroll: int | bool | None = None) -> Callable[Concatenate[CarryT, P], tuple[CarryT, OutputT_co]]

scan_via(fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None) -> Callable[[CarryT], tuple[CarryT, OutputT_co]]

Return a function that scans over the sequence using fn.

fn should take a block and a carry and return (carry, output). Semantics match :func:haliax.scan over the block axis.

`vmap_via(fn: Callable[..., OutputT_co]) -> Callable[..., OutputT_co]` ¤

vmap_via(fn: VmapFunction[M, P, OutputT_co]) -> Callable[P, OutputT_co]

vmap_via(fn: Callable[[M], OutputT_co]) -> Callable[[], OutputT_co]

Return a function that applies each block independently using fn.

fn should take a block and a carry and return an output. The returned function mirrors :func:haliax.vmap over the block axis.

`unstacked() -> Sequence[M]` ¤

`from_state_dict(state_dict: StateDict, prefix: str | None = None) -> M` ¤

`to_state_dict(prefix: str | None = 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.".

`get_layer(index: int) -> M` ¤

Return the indexth block in this sequential container.

`BlockFoldable` ¤

Bases: Protocol[M]

Common interface for :class:~haliax.nn.Stacked and :class:~haliax.nn.BlockSeq.

The interface exposes the :py:meth:fold and :py:meth:scan methods along with the helper methods :py:meth:fold_via, :py:meth:scan_via, and :py:meth:vmap_via.

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 Protocol classes, but you can use it as a type hint.

Methods:

init –
scan –
fold –
fold_via –
scan_via –
vmap_via –
unstacked –

Returns the unstacked version of this module. This is useful for logging or saving checkpoints.
get_layer –

Return the indexth layer of the folded module.

Attributes:

Block (Axis) –

`Block: Axis` `instance-attribute` ¤

`init(Block: Axis, module: Type[M], *, gradient_checkpointing: ScanCheckpointSpec = False, prevent_cse: bool = False) -> ModuleInit[S]` `classmethod` ¤

`scan(init: T, *extra_args, unroll: int | bool | None = None, **extra_kwargs)` ¤

`fold(init: T, *args, unroll: int | bool | None = None, **kwargs) -> T` ¤

`fold_via(fn: Callable[..., CarryT], *, unroll: int | bool | None = None) -> Callable[Concatenate[CarryT, P], CarryT]` ¤

fold_via(fn: FoldFunction[M, P, CarryT], *, unroll: int | bool | None = None) -> Callable[Concatenate[CarryT, P], CarryT]

fold_via(fn: Callable[[M, CarryT], CarryT], *, unroll: int | bool | None = None) -> Callable[[CarryT], CarryT]

`scan_via(fn: Callable[..., tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None) -> Callable[P, tuple[CarryT, OutputT_co]]` ¤

scan_via(fn: ScanFunction[M, CarryT, P, OutputT_co], *, unroll: int | bool | None = None) -> Callable[Concatenate[CarryT, P], tuple[CarryT, OutputT_co]]

scan_via(fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None) -> Callable[[CarryT], tuple[CarryT, OutputT_co]]

`vmap_via(fn: Callable[..., OutputT_co]) -> Callable[..., OutputT_co]` ¤

vmap_via(fn: VmapFunction[M, P, OutputT_co]) -> Callable[P, OutputT_co]

vmap_via(fn: Callable[[M], OutputT_co]) -> Callable[[], OutputT_co]

`unstacked() -> Sequence[M]` ¤

Returns the unstacked version of this module. This is useful for logging or saving checkpoints.

`get_layer(index: int) -> M` ¤

Return the indexth layer of the folded module.

Scan and Fold¤

haliax.scan¤

Key Features¤

Basic Example¤

Using scan with no inputs¤

haliax.fold¤

Basic Example¤

haliax.map¤

Gradient Checkpointing / Rematerialization¤

TL;DR Guidance¤

Simple Checkpointing¤

Nested Scan¤

Advanced: customizing the number of blocks¤

Custom Checkpointing Policies¤

Saving Block-Internal Values¤

Offloading Checkpointed Values¤

Summary of String and Boolean Aliases¤

Memory and Computation Tradeoffs¤

Module Stacks¤

Stacked¤

Creating a Stacked¤

Apply Blocks in Parallel with vmap¤

Fold Blocks vs Scan Blocks¤

Requirements for Stacked Blocks¤

BlockSeq and BlockFoldable¤

API¤

fold(fn: Callable, axis: AxisSelector, *, remat: ScanCheckpointSpec = False, reverse: bool = False, unroll: int = 1, is_scanned: BoolAxisSpec = is_named_or_shaped_array_like) -> Callable ¤

fn ¤

axis ¤

reverse ¤

unroll ¤

nested_scan ¤

is_scanned ¤

scan(f: Callable, axis: AxisSelector, *, remat: ScanCheckpointSpec = False, reverse=False, unroll=1, is_scanned: BoolAxisSpec = is_named_or_shaped_array_like) ¤

f ¤

axis ¤

reverse ¤

unroll ¤

remat ¤

is_scanned ¤

map(fn: Callable[[X], Y], axis: Axis, *, remat: ScanCheckpointSpec = False, reverse: bool = False, unroll: int = 1, is_mapped: BoolAxisSpec = is_jax_or_hax_array_like) -> Callable[[PyTree[X]], PyTree[Y]] ¤

fn ¤

axis ¤

reverse ¤

unroll ¤

Nested Remat¤

Offloading¤

save_carries: bool | Literal['offload'] = True class-attribute instance-attribute ¤

save_inputs: bool | Literal['offload'] = False class-attribute instance-attribute ¤

save_block_internals: bool | list[str] = False class-attribute instance-attribute ¤

offload_block_internals: list[str] = dataclasses.field(default_factory=list) class-attribute instance-attribute ¤

prevent_cse: bool = False class-attribute instance-attribute ¤

disable: bool = False class-attribute instance-attribute ¤

simple: bool = False class-attribute instance-attribute ¤

nested: bool | int = False class-attribute instance-attribute ¤

from_bool_or_str(remat_policy: bool | str) staticmethod ¤

checkpoint(carry_name: str, input_name: str, callable) ¤

Modules¤

Stacked ¤

stacked: M instance-attribute ¤

Block: Axis = eqx.field(static=True) class-attribute instance-attribute ¤

gradient_checkpointing: ScanCheckpointPolicy = eqx.field(static=True) class-attribute instance-attribute ¤

Layers: Axis property ¤

flatten_for_export() -> Mod ¤

unflatten_from_export(template: Mod) -> Mod ¤

init(Block: Axis, module: Type[M], *, gradient_checkpointing: ScanCheckpointSpec = False, prevent_cse: bool | None = None) -> ModuleInit[Stacked[M]] classmethod ¤

Block ¤

module ¤

gradient_checkpointing ¤

prevent_cse ¤

scan(init, *extra_args, unroll: int | bool | None = None, **extra_kwargs) ¤

init ¤

*extra_args ¤

**extra_kwargs ¤

fold(init, *args, unroll: int | bool | None = None, **kwargs) ¤

init ¤

*args ¤

**kwargs ¤

fold_via(fn: Callable[..., CarryT], *, unroll: int | bool | None = None) ¤

scan_via(fn: Callable[..., tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None) ¤

haliax.scan ¤

Using `scan` with no inputs¤

haliax.fold ¤

haliax.map ¤

Apply Blocks in Parallel with `vmap`¤

`fold(fn: Callable, axis: AxisSelector, *, remat: ScanCheckpointSpec = False, reverse: bool = False, unroll: int = 1, is_scanned: BoolAxisSpec = is_named_or_shaped_array_like) -> Callable` ¤

`fn` ¤

`axis` ¤

`reverse` ¤

`unroll` ¤

`nested_scan` ¤

`is_scanned` ¤

`scan(f: Callable, axis: AxisSelector, *, remat: ScanCheckpointSpec = False, reverse=False, unroll=1, is_scanned: BoolAxisSpec = is_named_or_shaped_array_like)` ¤

`f` ¤

`axis` ¤

`reverse` ¤

`unroll` ¤

`remat` ¤

`is_scanned` ¤

`map(fn: Callable[[X], Y], axis: Axis, *, remat: ScanCheckpointSpec = False, reverse: bool = False, unroll: int = 1, is_mapped: BoolAxisSpec = is_jax_or_hax_array_like) -> Callable[[PyTree[X]], PyTree[Y]]` ¤

`fn` ¤

`axis` ¤

`reverse` ¤

`unroll` ¤

`save_carries: bool | Literal['offload'] = True` `class-attribute` `instance-attribute` ¤

`save_inputs: bool | Literal['offload'] = False` `class-attribute` `instance-attribute` ¤

`save_block_internals: bool | list[str] = False` `class-attribute` `instance-attribute` ¤

`offload_block_internals: list[str] = dataclasses.field(default_factory=list)` `class-attribute` `instance-attribute` ¤

`prevent_cse: bool = False` `class-attribute` `instance-attribute` ¤

`disable: bool = False` `class-attribute` `instance-attribute` ¤

`simple: bool = False` `class-attribute` `instance-attribute` ¤

`nested: bool | int = False` `class-attribute` `instance-attribute` ¤

`from_bool_or_str(remat_policy: bool | str)` `staticmethod` ¤

`checkpoint(carry_name: str, input_name: str, callable)` ¤

`Stacked` ¤

`stacked: M` `instance-attribute` ¤

`Block: Axis = eqx.field(static=True)` `class-attribute` `instance-attribute` ¤

`gradient_checkpointing: ScanCheckpointPolicy = eqx.field(static=True)` `class-attribute` `instance-attribute` ¤

`Layers: Axis` `property` ¤

`flatten_for_export() -> Mod` ¤

`unflatten_from_export(template: Mod) -> Mod` ¤

`init(Block: Axis, module: Type[M], *, gradient_checkpointing: ScanCheckpointSpec = False, prevent_cse: bool | None = None) -> ModuleInit[Stacked[M]]` `classmethod` ¤

`Block` ¤

`module` ¤

`gradient_checkpointing` ¤

`prevent_cse` ¤

`scan(init, *extra_args, unroll: int | bool | None = None, **extra_kwargs)` ¤

`init` ¤

**`*extra_args`** ¤

`extra_kwargs`** ¤

`fold(init, *args, unroll: int | bool | None = None, **kwargs)` ¤

`init` ¤

**`*args`** ¤

`kwargs`** ¤

`fold_via(fn: Callable[..., CarryT], *, unroll: int | bool | None = None)` ¤

`scan_via(fn: Callable[..., tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None)` ¤

`vmap_via(fn: Callable[..., OutputT_co]) -> Callable[..., OutputT_co]` ¤

`vmap(*extra_args, **extra_kwargs)` ¤

`unstacked() -> Sequence[M]` ¤

`get_layer(index: int) -> M` ¤

`to_state_dict(prefix: str | None = None) -> StateDict` ¤

`from_state_dict(state_dict: StateDict, prefix: str | None = None) -> M` ¤

`BlockSeq` ¤

`blocks: Sequence[M]` `instance-attribute` ¤

`Block: Axis = eqx.field(static=True)` `class-attribute` `instance-attribute` ¤

`gradient_checkpointing: ScanCheckpointPolicy = eqx.field(static=True)` `class-attribute` `instance-attribute` ¤

`flatten_for_export() -> Mod` ¤

`unflatten_from_export(template: Mod) -> Mod` ¤

`init(Block: Axis, module: Type[M], *, gradient_checkpointing: ScanCheckpointSpec = False, prevent_cse: bool | None = None) -> ModuleInit[S]` `classmethod` ¤

`scan(init: T, *extra_args, unroll: int | bool | None = None, **extra_kwargs)` ¤

`fold(init: T, *args, unroll: int | bool | None = None, **kwargs) -> T` ¤

`fold_via(fn: Callable[..., CarryT], *, unroll: int | bool | None = None)` ¤

`scan_via(fn: Callable[..., tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None)` ¤

`vmap_via(fn: Callable[..., OutputT_co]) -> Callable[..., OutputT_co]` ¤

`unstacked() -> Sequence[M]` ¤

`from_state_dict(state_dict: StateDict, prefix: str | None = None) -> M` ¤

`to_state_dict(prefix: str | None = None) -> StateDict` ¤

`get_layer(index: int) -> M` ¤

`BlockFoldable` ¤

`Block: Axis` `instance-attribute` ¤