kups.core.effects ¶

class Interpreter[Context]:
    """Evaluate a jaxpr while threading a context through effect handlers.

    Args:
        handlers: Handler per primitive. Every equation of a listed primitive
            goes to its handler; all other primitives are bound unchanged.
        effects: Primitives that perform effects. Applies ``policy`` when one is
            not handled, directly or inside an unhandled primitive's sub-jaxpr.
        policy: What to do with such an unreachable effect.
        label: Name used in tracing debug info.
    """

    def __init__(
        self,
        handlers: Mapping[Primitive | str, Handler[Context]],
        *,
        effects: Collection[Primitive | str] = (),
        policy: InterpreterPolicy = InterpreterPolicy.RAISE,
        label: str = "",
    ):
        self.handlers: Mapping[str, Handler[Context]] = {
            _name(p): h for p, h in handlers.items()
        }
        self.effects: frozenset[str] = frozenset(map(_name, effects))
        self.policy = policy
        self.label = label

    def _has_effect(self, eqn: JaxprEqn) -> bool:
        """Whether ``eqn`` is an effect or has one in a sub-jaxpr."""
        return eqn.primitive.name in self.effects or any(
            self._has_effect(inner)
            for jaxpr in _sub_jaxprs(eqn)
            for inner in jaxpr.eqns
        )

    def _handler(self, eqn: JaxprEqn) -> Handler[Context] | None:
        handler = self.handlers.get(eqn.primitive.name)
        if handler is not None or not self._has_effect(eqn):
            return handler
        message = f"No handler for {eqn.primitive}, which performs or contains effects."
        if self.policy == InterpreterPolicy.RAISE:
            raise NotImplementedError(message)
        if self.policy == InterpreterPolicy.WARN:
            warnings.warn(f"{message} Its effects are ignored.", stacklevel=2)
        return None

    def __call__(
        self, jaxpr: ClosedJaxpr, ctx: Context, *args: Any
    ) -> tuple[list[TracerValue], Context]:
        env = InterpreterEnvironment(jaxpr.jaxpr, jaxpr.consts, args)
        for eqn in jaxpr.jaxpr.eqns:
            try:
                invals = env.read_many(eqn.invars)
            except Exception as e:
                raise ValueError(
                    f"Error reading from env: primitive={eqn.primitive}, invars={eqn.invars}"
                ) from e

            name_stack = (
                source_info_util.current_name_stack() + eqn.source_info.name_stack
            )
            with source_info_util.user_context(
                eqn.source_info.traceback, name_stack=name_stack
            ):
                handler = self._handler(eqn)
                if handler is None:
                    outvals = bind_primitive(eqn, invals)
                else:
                    ctx, outvals = handler(self, ctx, eqn, invals)

            try:
                env.write_many(eqn.outvars, outvals)
            except Exception as e:
                raise ValueError(
                    f"Error writing to env: primitive={eqn.primitive}, outvars={eqn.outvars} outvals={outvals}"
                ) from e

        try:
            out = env.read_many(jaxpr.jaxpr.outvars)
        except Exception as e:
            raise ValueError(f"Error reading from env: {jaxpr.jaxpr.outvars}") from e

        return out, ctx

class InterpreterPolicy(Enum):
    """What to do with an effect the interpreter cannot reach.

    That is an effect primitive without a handler, or an unhandled primitive
    whose sub-jaxprs contain an effect primitive.
    """

    RAISE = "raise"
    """Raise ``NotImplementedError``."""

    WARN = "warn"
    """Warn, then bind the primitive unchanged. Its effects are lost."""

    IGNORE = "ignore"
    """Bind the primitive unchanged. Its effects are lost."""