Actions

actions

Optional native state-derivative actions on compiled RHS and observables.

NativeStateActions(model, backend, observable, state_names, parameter_names, output_shape, _function, _transforms) dataclass

Caller-owned JAX/Torch state JVP and transpose/VJP adapters.

The transform acts directly on a compiled callable, without requesting a numerical Jacobian matrix. Forward/reverse AD storage and compiler memory remain backend/workload dependent. No state, bound parameter, reverse tape or solve cache survives a call. Parameters and time are fixed for the state action; separately qualified outer native differentiation can retain them.

Attributes:

Name	Type	Description
model	StructuralModelIdentity	Existing compiler source identity in natural state coordinates.
backend	Literal['jax', 'torch']	Explicit optional native AD backend.
observable	str | None	Named observable, or None for the RHS.
state_names	tuple[str, ...]	Ordered flat current-state coordinates.
parameter_names	tuple[str, ...]	Required runtime inputs for this target callable.
output_shape	tuple[int, ...]	RHS vector or declared observable shape.

jvp(t, y, tangent, **parameters)

Apply the target's state Jacobian to a flat native tangent.

Returns:

Type	Description
object	Native output with the target's declared output shape.

Raises:

Type	Description
ValueError	Input names, layout, precision or backend differ.

Source code in src/op_system/actions.py

def jvp(
    self, t: object, y: object, tangent: object, **parameters: object
) -> object:
    """Apply the target's state Jacobian to a flat native tangent.

    Returns:
        Native output with the target's declared output shape.

    Raises:
        ValueError: Input names, layout, precision or backend differ.
    """
    function = self._binding(t, y, parameters)
    self._native(tangent, y, "tangent")
    if getattr(tangent, "shape", None) != (len(self.state_names),):
        msg = "Tangent shape differs from the flat state layout"
        raise ValueError(msg)
    return self._transforms.jvp(function, (y,), (tangent,))[1]

vjp(t, y, cotangent, **parameters)

Apply the target's transpose state Jacobian to a native cotangent.

Returns:

Type	Description
object	Native flat state cotangent. The reverse pullback is call-local.

Raises:

Type	Description
ValueError	Input names, layout, precision or backend differ.

Source code in src/op_system/actions.py

def vjp(
    self, t: object, y: object, cotangent: object, **parameters: object
) -> object:
    """Apply the target's transpose state Jacobian to a native cotangent.

    Returns:
        Native flat state cotangent. The reverse pullback is call-local.

    Raises:
        ValueError: Input names, layout, precision or backend differ.
    """
    function = self._binding(t, y, parameters)
    self._native(cotangent, y, "cotangent")
    if getattr(cotangent, "shape", None) != self.output_shape:
        msg = "Cotangent shape differs from the target output layout"
        raise ValueError(msg)
    _, pullback = self._transforms.vjp(function, y)
    return pullback(cotangent)[0]

state_actions(compiled, *, backend, observable=None)

Bind optional native state actions without constructing a dense Jacobian.

Runtime inputs are homogeneous native float32/64 arrays. Plain real scalar time is also accepted. Parameter shapes retain the compiler's validation. Array observations use their original declared shape, not a flattened guess. The bounded dense symbolic API remains available separately for NumPy.

Parameters:

Name	Type	Description	Default
compiled	CompiledRhs	Unmodified compiler-produced source in natural coordinates.	required
backend	Literal['jax', 'torch']	JAX or Torch; imported only when explicitly requested.	required
observable	str | None	One compiled observable name, or None for the RHS.	None

Returns:

Type	Description
NativeStateActions	Static source/layout adapter with call-local native forward/reverse AD.

Raises:

Type	Description
ValueError	Backend, source or requested observable is invalid.
NotImplementedError	Operators/history/gauge-fixed execution is declared.
ImportError	The selected optional native backend is unavailable.

Source code in src/op_system/actions.py

def state_actions(
    compiled: CompiledRhs,
    *,
    backend: Literal["jax", "torch"],
    observable: str | None = None,
) -> NativeStateActions:
    """Bind optional native state actions without constructing a dense Jacobian.

    Runtime inputs are homogeneous native float32/64 arrays. Plain real scalar
    time is also accepted. Parameter shapes retain the compiler's validation.
    Array observations use their original declared shape, not a flattened guess.
    The bounded dense symbolic API remains available separately for NumPy.

    Args:
        compiled: Unmodified compiler-produced source in natural coordinates.
        backend: JAX or Torch; imported only when explicitly requested.
        observable: One compiled observable name, or None for the RHS.

    Returns:
        Static source/layout adapter with call-local native forward/reverse AD.

    Raises:
        ValueError: Backend, source or requested observable is invalid.
        NotImplementedError: Operators/history/gauge-fixed execution is declared.
        ImportError: The selected optional native backend is unavailable.
    """
    if backend not in {"jax", "torch"}:
        msg = "Native state actions require explicit jax or torch backend"
        raise ValueError(msg)
    rhs = compiled._rhs  # ruff: ignore[private-member-access]
    if rhs is None:
        msg = "Native state actions require retained compiler source"
        raise ValueError(msg)
    if tuple(rhs.state_names) != compiled.state_names:
        msg = "Native action source and compiled state layouts differ"
        raise ValueError(msg)
    if (
        compiled.operators
        or compiled.history_requirements
        or compiled.gauge_fix is not None
    ):
        msg = "Native actions currently exclude operators/history/gauge-fixed execution"
        raise NotImplementedError(msg)
    model = StructuralModelIdentity.from_compiled(compiled)
    function: _StateFunction = compiled.eval_fn
    shape: tuple[int, ...] = (len(compiled.state_names),)
    names = tuple(
        dict.fromkeys((
            *compiled.param_names,
            *(name for name, _ in rhs.shaped_params),
            *(name for name, _ in rhs.time_varying_params),
        ))
    )
    if observable is not None:
        if observable not in compiled.observables:
            msg = "Unknown compiled observable target"
            raise ValueError(msg)
        target = compiled.observables[observable]
        function = target.eval_fn
        shape = target.output_shape
        names = tuple(name for name, _ in target.parameter_shapes)
    try:
        module = import_module(backend)
    except ImportError as error:
        msg = f"Native actions require the optional op_system[{backend}] dependency"
        raise ImportError(msg) from error
    transforms = cast("_Transforms", module if backend == "jax" else module.func)
    return NativeStateActions(
        model,
        backend,
        observable,
        compiled.state_names,
        names,
        shape,
        function,
        transforms,
    )