kups.core.result ¶

@dataclass
class Result[State, Return]:
    """
    A result of a computation, containing the actual result and a sequence of runtime
    assertions that will be checked after the computation is done (outside of
    JIT-compiled functions).
    """

    value: Return
    assertions: tuple[RuntimeAssertion[State, Any], ...] = ()
    _all_pass: Array = field(default=None, kw_only=True)  # type: ignore

    def __iter__(self):
        yield self.value
        yield from self.assertions

    @skip_post_init_if_disabled
    def __post_init__(self):
        # Let's precompute _all_pass to compute its value during initialization.
        # This way we avoid computing it multiple times when checking assertions later.
        if self._all_pass is None:
            object.__setattr__(
                self,
                "_all_pass",
                jnp.all(jnp.array([a.predicate for a in self.assertions])),
            )

    @no_jax_tracing
    def raise_assertion(self):
        """Raises an assertion error if any of the runtime assertions failed.

        Note:
        This method should be called outside of JIT-compiled functions.
        """
        for assertion in self.failed_assertions:
            raise assertion.exception

    @property
    def all_assertions_pass(self) -> Array:
        """
        Returns True if all runtime assertions pass, i.e., their predicates
        evaluate to True. By compressing all predicates into a single boolean array,
        we only need to transfer one boolean value from device to host to check if
        all assertions passed.
        """
        return self._all_pass

    @property
    @no_jax_tracing
    def failed_assertions(self) -> tuple[RuntimeAssertion[State, Any], ...]:
        """
        Returns a tuple of runtime assertions that failed, i.e., those
        whose predicate evaluates to False.

        Note:
        This method should be called outside of JIT-compiled functions.
        """
        try:
            if self.all_assertions_pass.item():
                return ()
            return tuple(
                e for e in self.assertions if not bool(jnp.min(e.predicate).item())
            )
        except jax.errors.ConcretizationTypeError as e:
            logging.error(
                "Attempting to evaluate runtime assertions with non-concrete values."
                " You most likely attempted this within a jit-compiled function."
                " Check assertions outside of jit-compiled functions."
            )
            raise e

    @no_jax_tracing
    def fix_or_raise(self, state: State) -> State:
        """Apply fixes for all failed assertions, raising for any without a fix function.

        If all assertions pass, the state is returned unchanged. Otherwise each
        failed assertion's fix function is applied in sequence. Assertions that
        have no fix function registered will raise their configured exception.

        Args:
            state: Current simulation state to repair.

        Returns:
            State with all fixable assertion failures corrected.

        Raises:
            Exception: The configured exception of any failed assertion that has no fix function.
        """
        assertions = self.failed_assertions
        if not assertions:
            return state
        logging.info("Fixing failed assertions.")
        for assertion in assertions:
            logging.info("\t" + str(assertion.exception))
            state = assertion.fix(state)
        return state