kups.core.assertion.tracing ¶

def cond_handler(
    interpreter: Interpreter[AssertionContext],
    ctx: AssertionContext,
    eqn: JaxprEqn,
    invals: list[TracerValue],
) -> HandlerResult[AssertionContext]:
    """Custom cond handler that unions assertions across branches.

    Branches may declare different assertions. The output context holds the
    incoming assertions followed by every branch's assertions concatenated in
    branch order. Each branch fills its own slots with real predicates and the
    slots owned by other branches with passing placeholders (True predicate,
    zeroed fmt_args/fix_args), so jax.lax.cond observes matching pytree
    structures across branches. An unchosen branch thus contributes only passing
    assertions.
    """
    _, bind_params = get_bind_params(eqn)
    branches = bind_params["branches"]
    assert len(branches) > 0, "cond must have at least one branch"

    branch_fns = [
        jax.jit(reinterpret(jaxpr_as_fun(jaxpr), interpreter)) for jaxpr in branches
    ]
    # Incoming assertions bypass the conditional unless a branch reads them.
    needs_ctx = _contains_primitive(branches, check_assertion_p.name)
    base = ctx if needs_ctx else AssertionContext()
    n_in = len(base.assertions)

    # Dry-run trace each branch to learn the assertions it appends; the abstract
    # leaves (shape/dtype) are used to build placeholders for the other branches.
    suffix_templates = [
        fn.trace(base, *invals[1:]).out_info[1].assertions[n_in:] for fn in branch_fns
    ]
    packer = _AssertionPacker(sum(suffix_templates, ()))

    def wrap(
        index: int, fn: Callable[..., tuple[Any, AssertionContext]]
    ) -> Callable[..., tuple[Any, list[Array]]]:
        def wrapped(ctx_in: AssertionContext, *args: Any) -> tuple[Any, list[Array]]:
            with jax.disable_jit(False):
                outvals, ctx_out = fn(ctx_in, *args)
            merged: list[RuntimeAssertion[Any, Any]] = []
            for i, templates in enumerate(suffix_templates):
                merged.extend(
                    ctx_out.assertions[n_in:]
                    if i == index
                    else map(_passing, templates)
                )
            return outvals, packer.pack(tuple(merged))

        return wrapped

    wrapped_fns = [wrap(i, fn) for i, fn in enumerate(branch_fns)]

    # jax.lax.switch selects branch ``invals[0]`` (the cond/switch index) directly,
    # matching the primitive's branch order and supporting any number of branches.
    # Branches only append, so the incoming assertions need not be returned.
    outvals, packed = jax.lax.switch(invals[0], wrapped_fns, base, *invals[1:])
    return HandlerResult(
        AssertionContext(ctx.assertions + packer.unpack(packed)), outvals
    )

def scan_handler(
    interpreter: Interpreter[AssertionContext],
    ctx: AssertionContext,
    eqn: JaxprEqn,
    invals: list[TracerValue],
) -> HandlerResult[AssertionContext]:
    """Carry only the body's assertions through the scan.

    Each iteration folds the new assertions into the accumulated ones with their
    ``loop_merge``. Assertions created before the scan stay outside it. The body
    is traced once: the dry run that discovers its assertions is reused by the
    scan.
    """
    _, bind_params = get_bind_params(eqn)
    consts, carry, xs = split_scan_operands(bind_params, invals)
    n_in = len(ctx.assertions)
    body = jax.jit(partial(interpreter, bind_params["jaxpr"]))
    _, ctx_out_tree = body.trace(ctx, *consts, *carry, *(x[0] for x in xs)).out_info
    init = tuple(map(_initialize_loop_assertion, ctx_out_tree.assertions[n_in:]))

    def step(
        state: tuple[tuple[RuntimeAssertion[Any, Any], ...], list[Any]],
        x: list[Any],
    ) -> tuple[tuple[tuple[RuntimeAssertion[Any, Any], ...], list[Any]], list[Any]]:
        kept, args = state
        out, new_ctx = body(ctx, *consts, *args, *x)
        kept = tuple(
            _merge_loop_assertion(old, new)
            for old, new in zip(kept, new_ctx.assertions[n_in:], strict=True)
        )
        carry_out, ys = split_scan_results(bind_params, out)
        return (kept, carry_out), ys

    (kept, carry), ys = jax.lax.scan(
        step,
        (init, carry),
        xs,
        length=bind_params["length"],
        unroll=bind_params["unroll"],
        reverse=bind_params["reverse"],
    )
    return HandlerResult(AssertionContext(ctx.assertions + kept), [*carry, *ys])

def while_handler(
    interpreter: Interpreter[AssertionContext],
    ctx: AssertionContext,
    eqn: JaxprEqn,
    invals: list[TracerValue],
) -> HandlerResult[AssertionContext]:
    """Thread only the loop's own assertions, packed, through the carry.

    Each iteration folds the new assertions into the accumulated ones with their
    ``loop_merge``. Assertions created before the loop stay outside it.
    """
    _, bind_params = get_bind_params(eqn)
    cond_jaxpr: ClosedJaxpr = bind_params["cond_jaxpr"]
    body_jaxpr: ClosedJaxpr = bind_params["body_jaxpr"]
    num_cond_consts = bind_params["cond_nconsts"]
    n_consts = num_cond_consts + bind_params["body_nconsts"]
    cond_consts = invals[:num_cond_consts]
    body_consts = invals[num_cond_consts:n_consts]
    nonconst_invals = invals[n_consts:]
    n_in = len(ctx.assertions)
    body = jax.jit(partial(interpreter, body_jaxpr))
    _, ctx_out_tree = body.trace(ctx, *body_consts, *nonconst_invals).out_info
    templates = ctx_out_tree.assertions[n_in:]
    packer = _AssertionPacker(templates)

    def full_ctx(packed: list[Array]) -> AssertionContext:
        return AssertionContext(ctx.assertions + packer.unpack(packed))

    def _cond(carry: tuple[list[Array], list[Any]]) -> Array:
        packed, args = carry
        out, new_ctx = interpreter(cond_jaxpr, full_ctx(packed), *cond_consts, *args)
        if len(new_ctx.assertions) != n_in + len(templates):
            raise ValueError("While cond modified the context.")
        return jax.tree.leaves(out)[0]

    def _body(
        carry: tuple[list[Array], list[Any]],
    ) -> tuple[list[Array], list[Any]]:
        packed, args = carry
        out, new_ctx = body(ctx, *body_consts, *args)
        merged = map(
            _merge_loop_assertion, packer.unpack(packed), new_ctx.assertions[n_in:]
        )
        return packer.pack(tuple(merged)), list(out)

    init = packer.pack(tuple(map(_initialize_loop_assertion, templates)))
    packed, outvals = jax.lax.while_loop(_cond, _body, (init, list(nonconst_invals)))
    return HandlerResult(full_ctx(packed), outvals)

def with_runtime_assertions[**P, R](
    fn: Callable[P, R],
    policy: InterpreterPolicy = InterpreterPolicy.RAISE,
    context_sharding: jax.sharding.PartitionSpec | None = None,
) -> Callable[P, tuple[R, tuple[RuntimeAssertion[Any, Any], ...]]]:
    """
    Decorator that enables runtime assertion tracing for JAX functions.

    This decorator wraps a function to intercept and collect all runtime assertions
    created with `runtime_assert` during execution. The wrapped function returns
    both the original result and a tuple of all assertions that were evaluated,
    allowing for post-execution analysis, debugging, and optional error recovery.

    Args:
        fn: The function to wrap with assertion tracing capabilities
        policy: What to do with assertions the interpreter cannot reach, e.g.
            inside ``jax.custom_jvp``: RAISE (default), WARN or IGNORE (see
            [InterpreterPolicy][kups.core.effects.interpreter.InterpreterPolicy]).
        context_sharding: Optional sharding specification for distributed contexts.
            When provided, assertion contexts are sharded according to this spec
            for multi-device computations.

    Returns:
        A wrapped function that returns a tuple of (original_result, assertions_tuple).
        The assertions tuple contains all RuntimeAssertion instances encountered
        during execution, preserving order and enabling post-hoc analysis.

    Type Parameters:
        P: Parameter specification of the wrapped function (ParamSpec)
        R: Return type of the wrapped function

    Example:
        Basic usage:
        ```python
        @with_runtime_assertions
        def validate_computation(x):
            runtime_assert(x > 0, "x must be positive")
            return x ** 2

        result, assertions = validate_computation(jnp.array(5.0))
        # result = 25.0, assertions contains one RuntimeAssertion
        ```

        With custom policy:
        ```python
        traced_fn = with_runtime_assertions(
            my_function,
            policy=InterpreterPolicy.WARN
        )
        result, assertions = traced_fn(inputs)
        ```

        Distributed computation:
        ```python
        sharded_fn = with_runtime_assertions(
            distributed_computation,
            context_sharding=jax.sharding.PartitionSpec('data', None)
        )
        ```

        Error analysis and recovery:
        ```python
        result, assertions = traced_fn(initial_state)

        # Check for failures
        failed_assertions = [a for a in assertions if a.failed()]
        if failed_assertions:
            # Attempt automatic fixing
            fixed_state = initial_state
            for assertion in failed_assertions:
                if assertion.fix_fn is not None:
                    fixed_state = assertion.fix(fixed_state)

            # Re-run with fixed state
            result, _ = traced_fn(fixed_state)
        ```

    Note:
        The decorator integrates seamlessly with JAX transformations including
        jit, vmap, grad, and scan. Assertions are properly threaded through
        control flow operations and maintain correct semantics under
        transformations.
    """
    interpreter = Interpreter(
        {
            assertion_p: assertion_handler,
            check_assertion_p: check_assertion_handler,
            # Re-tracing jit and cond bodies yields slightly faster CPU programs.
            "jit": default_jit_handler,
            "cond": cond_handler,
            "remat2": default_checkpoint_handler,
            "scan": _only_with_assertions(scan_handler),
            "while": _only_with_assertions(while_handler),
            "shard_map": partial(shard_map_handler, context_sharding=context_sharding),
        },
        effects=(assertion_p, check_assertion_p),
        policy=policy,
        label="assertion_interpreter",
    )

    reinterpreted = reinterpret(fn, interpreter=interpreter)

    def wrapped(
        *args: P.args, **kwargs: P.kwargs
    ) -> tuple[R, tuple[RuntimeAssertion[Any, Any], ...]]:
        ctx = AssertionContext()
        outvals, ctx = reinterpreted(ctx, *args, **kwargs)
        return outvals, ctx.assertions

    return wrapped