kups.core.assertion ¶

def check_assertions(like: Array | None = None) -> Array:
    """A primitive that returns a scalar bool. When not wrapped in with_runtime_assertions,
    it always returns True. When wrapped, it will return the conjunction of all
    runtime assertions in the current context.

    Args:
        like: An array whose device placement and sharding will be used for the output.
            If None, the output will be placed on the default device.
    Returns:
        A scalar boolean array indicating whether all assertions pass.
    """
    if like is None:
        like = jnp.array(True, dtype=jnp.bool)
    return check_assertion_p.bind(like)

def runtime_assert[State, FixArgs](
    predicate: Array,
    message: str = "",
    fmt_args: dict[str, Array] | None = None,
    exception_type: type[Exception] = AssertionError,
    static_info: dict[str, Any] | None = None,
    fix_fn: Fix[State, FixArgs] | None = None,
    fix_args: FixArgs | _NO_ARGS = NO_ARGS,
    loop_merge: LoopMerge = ELEMENTWISE_MAX,
) -> None:
    """
    Create a runtime assertion that integrates with JAX transformations.

    This function creates assertions that can be traced through JAX transformations
    including JIT compilation, automatic differentiation, and vectorization. The
    assertion acts as an identity function during execution but records assertion
    metadata for later inspection.

    Args:
        predicate: A boolean array indicating whether the assertion passes
        message: Error message with optional format placeholders (e.g., "Value {val} is invalid")
        fmt_args: Dictionary mapping format placeholder names to values
        exception_type: Type of exception to raise if assertion fails during checking
        static_info: Additional metadata for debugging (not traced by JAX)
        fix_fn: Optional function to repair state when assertion fails
        fix_args: Arguments for the fix function (can be complex PyTree structures)
        loop_merge: How ``scan`` and ``while_loop`` fold ``fmt_args`` and
            ``fix_args`` across iterations. Defaults to
            [ELEMENTWISE_MAX][kups.core.assertion.ELEMENTWISE_MAX].

    Type Parameters:
        State: Type of state that can be modified by the fix function
        FixArgs: Type of arguments passed to the fix function (supports PyTree structures)

    Example:
        Basic assertion:
        ```python
        x = jnp.array(5.0)
        runtime_assert(
            predicate=x > 0,
            message="Value must be positive, got {val}",
            fmt_args={"val": x}
        )
        ```

        Assertion with fixing:
        ```python
        runtime_assert(
            predicate=x > threshold,
            message="Value {val} below threshold {thresh}",
            fmt_args={"val": x, "thresh": threshold},
            fix_fn=lambda state, args: jnp.maximum(state, args["min_val"]),
            fix_args={"min_val": jnp.array(1.0)}
        )
        ```

        Complex PyTree fix_args:
        ```python
        complex_args = {
            "thresholds": {"min": jnp.array(0.1), "max": jnp.array(10.0)},
            "multipliers": (jnp.array(2.0), jnp.array(3.0))
        }
        runtime_assert(
            predicate=x > 0,
            message="Invalid value",
            fix_args=complex_args
        )
        ```

    Note:
        The fix_args parameter supports arbitrarily nested PyTree structures including
        dictionaries, tuples, and arrays. These are automatically flattened and
        unflattened during JAX transformations.
    """
    if fmt_args is None:
        fmt_args = {}
    if static_info is None:
        static_info = {}

    tb = _capture_traceback().replace("{", "{{").replace("}", "}}")
    # Convert static_info to hashable format (tuple of key-value pairs)
    static_info_hashable = tuple(sorted(static_info.items()))

    # Prepare inputs: predicate, fmt_args values, fix_args flattened (if present)
    inputs = [predicate, *fmt_args.values()]
    fix_args_tree = None
    if not isinstance(fix_args, _NO_ARGS):
        # Flatten fix_args PyTree and add to inputs
        fix_args_flat, fix_args_tree = jax.tree.flatten(fix_args)
        inputs.extend(fix_args_flat)

    assertion_p.bind(
        *inputs,
        fmt_arg_names=tuple(fmt_args.keys()),
        message=message + f"{_TRACEBACK_MARKER}{tb}",
        exception_type=exception_type,
        static_info_hashable=static_info_hashable,
        fix_fn=fix_fn,
        fix_args_tree=fix_args_tree,
        loop_merge=loop_merge,
    )

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