Function scan_in_context

axlearn/common/module.py:1270–1377

A thin wrapper around `jax.lax.scan` which is compatible with `OutputCollection`. In particular, summaries and outputs added by `add_summary` and `add_module_output` respectively are accumulated in `current_context().output_collection`, subject to any output filtering. Specifically, sum

(
    fn,
    *,
    carry: NestedTensor,
    xs: NestedTensor,
    drop_output: Optional[Callable[[str], bool]] = None,
    child_name_prefix: str = "iter",
    unroll: Union[int, bool] = 1,
    remat_kwargs: Optional[dict[str, Any]] = None,
    merge_summaries: bool = False,
)

    fn = _Functional(
        context=context, method_fn=method_fn, require_parent=True, copy_args_tree=copy_args_tree
    )
    method_outputs, output_collection = fn(*args, **kwargs)

    for output_collection_type in drop_output_collections:
        getattr(output_collection, output_collection_type).clear()
    return method_outputs, output_collection


def scan_in_context(
    fn,
    *,
    carry: NestedTensor,
    xs: NestedTensor,
    drop_output: Optional[Callable[[str], bool]] = None,
    child_name_prefix: str = "iter",
    unroll: Union[int, bool] = 1,
    remat_kwargs: Optional[dict[str, Any]] = None,
    merge_summaries: bool = False,
) -> tuple[NestedTensor, NestedTensor]:
    """A thin wrapper around `jax.lax.scan` which is compatible with `OutputCollection`.

    In particular, summaries and outputs added by `add_summary` and `add_module_output` respectively
    are accumulated in `current_context().output_collection`, subject to any output filtering.
    Specifically, summaries from iteration `i` will be placed in
    `summaries[f"{child_name_prefix}{i}"]`; module outputs will be stacked and placed in
    `module_outputs[child_name_prefix]`.

    Args:
        fn: A function with args (carry, x) returning a dict(carry=..., y=...).
        carry: The initial value of the loop carry, to be accumulated across scan.
        xs: A dict with at least "x" as a key, where each leaf is a tensor of shape
            [num_scan_iters, ...]. At scan iteration i:
            - xs["x"][i, ...] represents the inputs to `fn`.
            - xs[key][i, ...] is provided as a kwarg to the ith invocation context.
        drop_output: A callable that takes a path and outputs a decision of whether to drop the
            output at the given path, where True means we drop. By default, the callable is None,
            meaning nothing is dropped.
        child_name_prefix: The child name prefix used for children to be added to
            `target_output_collection`.
        unroll: If a positive integer is provided, it determines how many unrolled loop iterations
            to run within a single rolled iteration of the loop. If a boolean is provided, it will
            determine if the loop is competely unrolled (i.e. unroll=True) or left completely rolled
            (i.e. unroll=False).
        remat_kwargs: Optional dict passed to `jax.checkpoint` to enable rematerialization.
            Common options include:
              - `prevent_cse`: (bool) Whether to disable common subexpression elimination.
                    If left unspecified, defaults to False following recommendations from the JAX
                    documentation.
                    Raises a ValueError if `prevent_cse` is set to True.
              - `policy`: A checkpoint policy from `jax.checkpoint_policies`.
            If provided, the scan body will be wrapped as:
                `scan_fn = jax.checkpoint(scan_fn, **remat_kwargs)`
            Otherwise, `jax.checkpoint` is not used.
            See https://docs.jax.dev/en/latest/_autosummary/jax.checkpoint.html.
        merge_summaries: If True, accumulate summaries across scan iterations via
            Summary.accumulate() instead of creating per-iteration children.