kups.core.effects.handlers ¶

def default_shard_map_handler[Context](
    interpreter: Interpreter[Context],
    ctx: Context,
    eqn: JaxprEqn,
    invals: list[TracerValue],
    *,
    declare_ctx_in_specs: Callable[[Context], Context | PartitionSpec] | None = None,
    declare_ctx_out_specs: Callable[[Context], Context | PartitionSpec] | None = None,
    reduce_ctx: Callable[[Context], Context] | None = None,
) -> HandlerResult[Context]:
    """Reinterpret a ``shard_map`` equation through ``interpreter``.

    Args:
        interpreter: The interpreter walking the surrounding jaxpr.
        ctx: Input context (lives outside the manual axis scope, typically
            replicated across all shards).
        eqn: The ``shard_map`` equation being handled.
        invals: User-data inputs to the equation.
        declare_ctx_in_specs: Optional. Map the input context tree to the
            ``in_specs`` shard_map should use for it. Defaults to ``P()``
            (replicated) for every leaf.
        declare_ctx_out_specs: Optional. Map the post-body context tree to
            the ``out_specs`` shard_map should use for it. Defaults to ``P()``
            (replicated, broadcast across leaves). Note: deriving the tree
            shape requires re-tracing the body outside the manual axis scope,
            which fails for bodies that bind manual axes (collectives, pvary,
            axis_index). Use ``reduce_ctx`` instead in that case.
        reduce_ctx: Optional. Transform applied to the per-shard post-body
            context inside the ``shard_map`` body before returning. The
            natural place to invoke collectives (``all_gather``, ``psum``, …)
            to combine per-shard data into globally-replicated values, since
            it executes inside the manual axis scope. When the result is
            fully replicated, the default ``out_specs=P()`` is satisfied
            without ``declare_ctx_out_specs``.
    """
    jaxpr = ClosedJaxpr(eqn.params["jaxpr"], ())

    def fn(ctx: Context, *invals: TracerValue) -> tuple[list[Any], Context]:
        outvals, ctx_post = interpreter(jaxpr, ctx, *invals)
        if reduce_ctx is not None:
            ctx_post = reduce_ctx(ctx_post)
        return outvals, ctx_post

    if declare_ctx_in_specs is not None:
        ctx_in_specs = declare_ctx_in_specs(jax.tree.map(Uninitialized, ctx))
    else:
        ctx_in_specs = jax.P()

    if declare_ctx_out_specs is not None:
        abstract_invals = [Uninitialized(x) for x in jaxpr.in_avals]
        ctx_out_tree = (
            jax.jit(fn)
            .trace(jax.tree.map(Uninitialized, ctx), *abstract_invals)
            .out_info[1]
        )
        ctx_out_specs = declare_ctx_out_specs(ctx_out_tree)
    else:
        ctx_out_specs = jax.P()

    sharded_fn = jax.shard_map(
        out_specs=(list(eqn.params["out_specs"]), ctx_out_specs),
        in_specs=(ctx_in_specs, *eqn.params["in_specs"]),
        mesh=eqn.params["mesh"],
        axis_names=shard_map_manual_axes(eqn.params),
        check_vma=eqn.params["check_vma"],
    )(fn)

    outvals, ctx_out = sharded_fn(ctx, *invals)
    return HandlerResult(ctx_out, outvals)