OP System

op_system

op_system.

Domain-agnostic RHS specification + compilation utilities.

Public API (v1)

Primary user entrypoints: - compile_spec: Validate, normalize, and compile a RHS specification in one step. - normalize_rhs: Validate and normalize a YAML-friendly RHS specification. - compile_rhs: Compile a NormalizedRhs into an efficient callable RHS.

Core data structures: - NormalizedRhs - CompiledRhs - OperatorDescriptor

Design guarantees: - No dependency on provider/adapters (eg flepimop2). - Stable interface for downstream engines. - Forward-compatible with multiphysics extensions.

IdentifierString = Annotated[str, AfterValidator(_validate_identifier_string)] module-attribute

Custom pydantic type for validated identifier strings used in op_system.

Identifier strings are used for state names, dimension names, and other keys in the system. They must be non-empty, contain only alphanumeric characters, and start with a letter. Leading and trailing whitespace is stripped before validation.

Examples:

>>> from pydantic import BaseModel
>>> from op_system import IdentifierString
>>> class ExampleModel(BaseModel):
...     identifier: IdentifierString
...
>>> ExampleModel(identifier="S")
ExampleModel(identifier='S')
>>> ExampleModel(identifier="  Foobar  ")
ExampleModel(identifier='Foobar')
>>> ExampleModel(identifier="123abc")
Traceback (most recent call last):
    ...
pydantic_core._pydantic_core.ValidationError: 1 validation error for ExampleModel
identifier
Value error, IdentifierString must contain only alphanumerical characters and start with a letter. [...]
    For further information visit ...
>>> ExampleModel(identifier="")
Traceback (most recent call last):
    ...
pydantic_core._pydantic_core.ValidationError: 1 validation error for ExampleModel
identifier
Value error, IdentifierString must not be empty. [...]
    For further information visit ...

Array

Bases: Protocol

Structural Array-API protocol.

Any object whose runtime type implements shape, dtype, __array_namespace__ and item satisfies this protocol. NumPy

= 2.0 ndarrays and JAX arrays (concrete and traced) qualify directly.

Runtime evaluation is deliberately broader: it discovers namespaces via :func:array_api_compat.array_namespace, so native arrays such as torch.Tensor are accepted even though they do not structurally satisfy this protocol. No compile-time backend selector is needed.

BlockAxisInfo(name, size, state_axis_pos, param_axis_pos) dataclass

Metadata for one block-diagonal (factorizable) axis.

A BlockAxisInfo is attached to :class:~op_system.compile.CompiledRhs for each axis declared in spec["factorize_axes"] that passes the IR separability check in :func:analyze_block_axes. Engines (e.g. the diffrax plugin) consume this to partition ODE solves with jax.vmap over the block axis.

Attributes:

Name	Type	Description
name	str	Axis name string (e.g. "loc").
size	int	Number of elements along the axis.
state_axis_pos	dict[str, int]	Maps each state-template base name to the integer position of this axis within that template's shape tuple. Only templates that carry this axis appear in the dict.
param_axis_pos	dict[str, int | None]	Maps each shaped parameter name to the integer position of this axis within the actual runtime array that the engine passes to the eval function, or None if the parameter does not carry this axis (broadcast). For non-time-varying shaped parameters the position is the index in the parameter's axis tuple. For time-varying parameters the runtime array has time prepended at index 0, so the position equals the index of the block axis in the full (time, *spatial_axes) tuple. Parameters that are entirely scalar (not shaped) do not appear in this dict and are always broadcast.

Note

BlockAxisInfo uses dict fields and therefore cannot be used as a hash key. It is stored on :class:~op_system.compile.CompiledRhs with hash=False and is fully pickle-stable.

Examples:

>>> info = BlockAxisInfo(
...     name="loc",
...     size=3,
...     state_axis_pos={"S": 1, "I": 1, "R": 1},
...     param_axis_pos={"rho": 0, "beta": None},
... )
>>> info.name
'loc'
>>> info.size
3
>>> info.state_axis_pos["S"]
1
>>> info.param_axis_pos["rho"]
0
>>> info.param_axis_pos["beta"] is None
True

BodyEvalFn

Bases: Protocol

Callable that evaluates history signal bodies at (t, y, **params).

Returns a mapping from signal_id to the evaluated body array/value.

CompiledReactant(state_base, state_axes, full_axes, pinned, order) dataclass

Molecular reactant metadata aligned to an expanded reaction channel.

The record is numerical-array neutral. Providers combine state_axes with the enclosing reaction's channel coordinates and pinned indices to locate a state cell, then place order in their reactant-order matrix. Catalysts appear here even when their net stoichiometric change is zero.

CompiledReaction(name, from_base, from_axes, full_axes, to_base, to_axes, sum_axes, pinned, from_pinned, propensity_fn, reactants=(), reactants_complete=False) dataclass

Compiled propensity + axis bookkeeping for one named transition.

Produced from :class:op_system._reactions.ReactionArtifactIR at compile time. A firing event at from-cell index i (in the natural N-D shape given by from_axes) depletes 1 from from_base at that cell and adds 1 to to_base at the cell obtained by: keeping every axis in to_axes at i's value for that axis, and using the fixed coordinate index from pinned for every axis in pinned (this includes both axes that are wildcard on from_axes and pinned on to -- a "collapse to a fixed target" transition, where multiple source cells share one destination and the axis also appears in sum_axes -- and axes pinned on BOTH from and to, a "point-to-point" shift between two specific coordinates on the same axis, e.g. a vaccination-dose-progression transition, which do NOT appear in sum_axes since there's only ever one source value on that axis). Consumers that need a single combined scatter/reduce target (e.g. an engine applying many simultaneous firings) sum over sum_axes when depositing into to_base -- op_system does not do that summation itself, since whether/how to combine simultaneous firings across a summed axis is an execution-semantics decision for the consumer, not a compile-time one.

from_base is None for a SOURCE-ONLY (from: null) reaction -- an exogenous hazard with no compartment to deplete (e.g. cross-district case importation). A consumer must special-case this: skip every from-side depletion/clamp step and apply only the to_base scatter, using from_axes/full_axes (which are the to-side template's own axes in this case, not a nonexistent from-side one -- see ReactionArtifactIR).

CompiledRhs(state_names, param_names, eval_fn, meta=(lambda: MappingProxyType({}))(), operators=tuple(), factorize_axes=tuple(), block_axes=tuple(), pytree_eval_fn=None, template_shapes=None, block_pytree_eval_fn=None, block_template_shapes=None, history_requirements=tuple(), history_eval_fn=None, body_eval_fn=None, block_history_eval_fn=None, block_body_eval_fn=None, reactions=tuple(), _rhs=None) dataclass

Container for a compiled RHS evaluation function.

Instances produced by :func:compile_rhs retain a private reference to their source :class:NormalizedRhs so the container can be pickled and re-hydrated by re-running the compile pipeline on load. eval_fn itself is a closure (and on the vectorized path captures compiled code objects), so it is dropped from the pickle and rebuilt by :func:compile_rhs in :meth:__setstate__. Round-tripping a CompiledRhs therefore costs one compile on load and yields a functionally equivalent instance whose eval_fn produces identical outputs for identical inputs.

__getstate__()

Return picklable state.

The compiled eval_fn is a closure (and on the vectorized path captures compiled :class:types.CodeType objects), which is not portably picklable. Instead we serialize just the source :class:NormalizedRhs and let :meth:__setstate__ recompile.

Raises:

Type	Description
TypeError	If the source NormalizedRhs was not retained (i.e. the instance was constructed directly rather than via :func:compile_rhs).

Source code in src/op_system/compile.py

def __getstate__(self) -> dict[str, Any]:
    """Return picklable state.

    The compiled ``eval_fn`` is a closure (and on the vectorized path
    captures compiled :class:`types.CodeType` objects), which is not
    portably picklable. Instead we serialize just the source
    :class:`NormalizedRhs` and let :meth:`__setstate__` recompile.

    Raises:
        TypeError: If the source ``NormalizedRhs`` was not retained
            (i.e. the instance was constructed directly rather than
            via :func:`compile_rhs`).
    """
    if self._rhs is None:
        msg = (
            "CompiledRhs is not picklable: the source NormalizedRhs was "
            "not retained. Construct via compile_rhs() to produce a "
            "picklable CompiledRhs."
        )
        raise TypeError(msg)
    return {"_rhs": self._rhs}

__setstate__(state)

Restore by recompiling from the pickled :class:NormalizedRhs.

Source code in src/op_system/compile.py

def __setstate__(self, state: Mapping[str, Any]) -> None:
    """Restore by recompiling from the pickled :class:`NormalizedRhs`."""
    rhs = state["_rhs"]
    rebuilt = compile_rhs(rhs)
    # frozen+slots dataclass: bypass __setattr__ via object.__setattr__
    object.__setattr__(self, "state_names", rebuilt.state_names)
    object.__setattr__(self, "param_names", rebuilt.param_names)
    object.__setattr__(self, "eval_fn", rebuilt.eval_fn)
    object.__setattr__(self, "meta", rebuilt.meta)
    object.__setattr__(self, "operators", rebuilt.operators)
    object.__setattr__(self, "factorize_axes", rebuilt.factorize_axes)
    object.__setattr__(self, "block_axes", rebuilt.block_axes)
    object.__setattr__(self, "pytree_eval_fn", rebuilt.pytree_eval_fn)
    object.__setattr__(self, "template_shapes", rebuilt.template_shapes)
    object.__setattr__(self, "block_pytree_eval_fn", rebuilt.block_pytree_eval_fn)
    object.__setattr__(self, "block_template_shapes", rebuilt.block_template_shapes)
    object.__setattr__(self, "history_requirements", rebuilt.history_requirements)
    object.__setattr__(self, "history_eval_fn", rebuilt.history_eval_fn)
    object.__setattr__(self, "body_eval_fn", rebuilt.body_eval_fn)
    object.__setattr__(self, "block_history_eval_fn", rebuilt.block_history_eval_fn)
    object.__setattr__(self, "block_body_eval_fn", rebuilt.block_body_eval_fn)
    object.__setattr__(self, "reactions", rebuilt.reactions)
    object.__setattr__(self, "_rhs", rhs)

bind(params)

Bind parameter values and return a 2-arg RHS: rhs(t, y) -> dydt.

Parameters:

Name	Type	Description	Default
params	Mapping[str, object]	Mapping of parameter names to values.	required

Returns:

Type	Description
Callable[[object, object], Float64Array]	A callable rhs(t, y) that evaluates the RHS with params fixed.

Source code in src/op_system/compile.py

def bind(
    self, params: Mapping[str, object]
) -> Callable[[object, object], Float64Array]:
    """Bind parameter values and return a 2-arg RHS: rhs(t, y) -> dydt.

    Args:
        params: Mapping of parameter names to values.

    Returns:
        A callable `rhs(t, y)` that evaluates the RHS with `params` fixed.
    """
    params_dict = dict(params)

    def rhs(t: object, y: object) -> Float64Array:
        return self.eval_fn(t, y, **params_dict)

    return rhs

EvalFn

Bases: Protocol

Callable RHS evaluator supporting runtime parameter kwargs.

Accepts a flat (n_state,) state array and returns a flat (n_state,) derivative array in the same array namespace.

ExprRhs(state_names, equations, aliases, param_names, all_symbols, meta, state_templates=(), shaped_params=(), time_varying_params=(), aliases_ir=dict(), equations_ir=(), aliases_ir_reduce=dict(), equations_ir_reduce=(), alias_templates=()) dataclass

Bases: _RhsBase

Normalized RHS for kind="expr" specs (explicit d(state)/dt equations).

Produced by :func:normalize_expr_rhs. Use :data:NormalizedRhs as the union type when you need to accept both kinds.

ExpressionString(source) dataclass

Validated expression wrapper with cached AST and symbol names.

as_ir()

Return the typed IR representation for this expression.

Returns:

Type	Description
Expr	Parsed typed IR tree.

Source code in src/op_system/_symbols.py

def as_ir(self) -> Expr:
    """Return the typed IR representation for this expression.

    Returns:
        Parsed typed IR tree.
    """
    return to_ir(self.ast)

as_lowered_ir()

Return helper-lowered typed IR for this expression.

Returns:

Type	Description
Expr	Typed IR tree with helper calls lowered to Reduce nodes.

Source code in src/op_system/_symbols.py

def as_lowered_ir(self) -> Expr:
    """Return helper-lowered typed IR for this expression.

    Returns:
        Typed IR tree with helper calls lowered to ``Reduce`` nodes.
    """
    return parse_expr_to_ir(self.source, lower_helpers=True)

OperatorDescriptor(axis, kind=None, bc=None, velocity=None, rate=None, kernel=None, name=None, apply_to=None, direction=None) dataclass

Typed description of a spatial operator declared in an op_system RHS spec.

Captures the model-level description of an operator that is known at compile time: which axis it acts on, its kind (advection, diffusion, etc.), optional boundary condition, and the names of any runtime parameters (velocity, rate) it consumes. Grid geometry, CN matrix construction, and solver staging remain downstream concerns handled by the engine.

Attributes:

Name	Type	Description
axis	str	Name of the axis the operator acts on (e.g. "loc").
kind	str | None	Operator type string, e.g. "advection", "diffusion", "transport", "jump_integral", or None if unspecified.
bc	str | None	Boundary condition, e.g. "absorbing", "periodic", "neumann", "reflecting", or None if unspecified.
velocity	str | float | None	Parameter name or numeric constant for an advection velocity, or None.
rate	str | float | None	Parameter name or numeric constant for a rate/coefficient, or None.
kernel	Mapping[str, Any] | None	Mixing-kernel sub-specification, or None.
name	str | None	Optional operator name from the specification.
apply_to	tuple[str, ...] | None	Concrete state names selected by the operator, or None when the operator applies to every compatible state.
direction	str | None	Optional normalized direction. For advection and transport, "increasing" preserves the resolved velocity sign and "decreasing" reverses it; when omitted, velocity is used as a signed coefficient. For jump integrals, "up", "down", or "both" masks source-to-target movement in coordinate order.

Examples:

>>> od = OperatorDescriptor(axis="loc")
>>> od.axis
'loc'
>>> od.kind is None
True
>>> od.bc is None
True
>>> od.velocity is None
True
>>> od2 = OperatorDescriptor(
...     axis="loc",
...     kind="advection",
...     bc="absorbing",
...     velocity=0.25,
...     name="drift",
...     apply_to=("S",),
... )
>>> od2.kind
'advection'
>>> od2.bc
'absorbing'
>>> od2.velocity
0.25
>>> od2.name
'drift'
>>> od2.apply_to
('S',)

PytreeEvalFn

Bases: Protocol

Callable RHS evaluator operating on shaped PyTree state dicts.

Accepts y as a StateDict (mapping from state-template base name to a shaped array with the template's natural N-D shape) and returns a StateDict of the same structure containing the derivative. Enables the engine to skip the flatten/unflatten step entirely and expose the full tensor structure to JAX/XLA.

ReactionPropensityFn

Bases: Protocol

Callable propensity evaluator for one named reaction/transition.

Accepts y as a StateDict and returns an array shaped like the reaction's from_axes (one independent rate per source cell) -- see :class:CompiledReaction.

ShapeGroup(cells, example_cell, example_expression) dataclass

Cells of one state template that share an expression shape.

StateString

Bases: BaseModel

Structured representation of a state string.

A state string is either a bare state name like "S" or a state name followed immediately by bracketed dimensions like "R[age,vax]".

Examples:

>>> StateString.model_validate("S")
StateString(name='S', dims=())
>>> recovery = StateString.model_validate("R[age,vax]")
>>> recovery
StateString(name='R', dims=('age', 'vax'))
>>> print(recovery)
R[age,vax]
>>> recovery.model_dump()
'R[age,vax]'
>>> StateString.model_validate("Foobar[ age , vax ]")
StateString(name='Foobar', dims=('age', 'vax'))

__str__()

Return the compact string form.

Returns:

Type	Description
str	Compact state string.

Examples:

>>> str(StateString(name="S", dims=()))
'S'
>>> str(StateString(name="R", dims=("age",)))
'R[age]'
>>> str(StateString(name="R", dims=("age", "vax")))
'R[age,vax]'
>>> str(StateString(name="lambda", dims=("age", "vax", "state")))
'lambda[age,vax,state]'

Source code in src/op_system/_state_string.py

def __str__(self) -> str:
    """
    Return the compact string form.

    Returns:
        Compact state string.

    Examples:
        >>> str(StateString(name="S", dims=()))
        'S'
        >>> str(StateString(name="R", dims=("age",)))
        'R[age]'
        >>> str(StateString(name="R", dims=("age", "vax")))
        'R[age,vax]'
        >>> str(StateString(name="lambda", dims=("age", "vax", "state")))
        'lambda[age,vax,state]'
    """
    if not self.dims:
        return self.name
    dims = ",".join(self.dims)
    return f"{self.name}[{dims}]"

TransitionsRhs(state_names, equations, aliases, param_names, all_symbols, meta, state_templates=(), shaped_params=(), time_varying_params=(), aliases_ir=dict(), equations_ir=(), aliases_ir_reduce=dict(), equations_ir_reduce=(), alias_templates=(), reactions_ir=()) dataclass

Bases: _RhsBase

Normalized RHS for kind="transitions" specs (per-capita hazard diagram).

Produced by :func:normalize_transitions_rhs. Use :data:NormalizedRhs as the union type when you need to accept both kinds.

ValidationReport(stages, errors=list(), cost=dict(), parameters=dict(), shape_groups=dict()) dataclass

Result of validating one op_system specification.

Attributes:

Name	Type	Description
stages	dict[str, str]	"passed", "failed", or "skipped" for normalize, compile, and vectorize.
errors	list[str]	Error messages from the failed stage.
cost	dict[str, int]	Compile-cost drivers (expanded states, templates, transitions, coordinate-pinned transitions, operators).
parameters	dict[str, tuple[str, ...]]	Parameters the spec consumes, mapped to their declared axes (empty tuple for scalars); operator parameters use kernel.param_axes when declared.
shape_groups	dict[str, list[ShapeGroup]]	When vectorization fails, the templates whose cells have more than one expression shape, with each distinct shape's count and an example cell. Empty otherwise, because comparing every cell's expanded equation is slow on large specs.

ok property

Whether every stage passed.

array_namespace(value)

Return an Array-API-compatible namespace for value.

array-api-compat recognizes both arrays that implement __array_namespace__ and native arrays such as torch.Tensor that need a compatibility namespace. Optional backends are imported only when one of their arrays is supplied.

Parameters:

Name	Type	Description	Default
value	object	Numerical array whose namespace should be selected.	required

Returns:

Type	Description
Any	The Array-API-compatible namespace for value.

Raises:

Type	Description
TypeError	If value is not a supported array object.

Source code in src/op_system/_array.py

def array_namespace(value: object) -> Any:  # ruff: ignore[any-type]
    """Return an Array-API-compatible namespace for ``value``.

    ``array-api-compat`` recognizes both arrays that implement
    ``__array_namespace__`` and native arrays such as ``torch.Tensor`` that
    need a compatibility namespace. Optional backends are imported only when
    one of their arrays is supplied.

    Args:
        value: Numerical array whose namespace should be selected.

    Returns:
        The Array-API-compatible namespace for ``value``.

    Raises:
        TypeError: If ``value`` is not a supported array object.
    """
    try:
        return _array_namespace(value)
    except TypeError as error:
        msg = (
            "op_system numerical inputs must be supported array objects; "
            f"got {type(value).__name__}."
        )
        raise TypeError(msg) from error

axis_kernel_generator_rhs(state, generator, *, axis, velocity=1.0)

Return velocity * state @ generator applied along axis.

Parameters:

Name	Type	Description	Default
state	Any	Array with the kernel axis at position axis.	required
generator	Any	Square matrix over the kernel axis; rows are sources.	required
axis	int	Position of the kernel axis in state.	required
velocity	Any	Scalar multiplier (for example a waning rate).	1.0

Returns:

Type	Description
Any	The derivative contribution, with the same shape as state.

Source code in src/op_system/_axis_kernel.py

def axis_kernel_generator_rhs(
    state: Any,  # ruff: ignore[any-type]
    generator: Any,  # ruff: ignore[any-type]
    *,
    axis: int,
    velocity: Any = 1.0,  # ruff: ignore[any-type]
) -> Any:  # ruff: ignore[any-type]
    """Return ``velocity * state @ generator`` applied along ``axis``.

    Args:
        state: Array with the kernel axis at position ``axis``.
        generator: Square matrix over the kernel axis; rows are sources.
        axis: Position of the kernel axis in ``state``.
        velocity: Scalar multiplier (for example a waning rate).

    Returns:
        The derivative contribution, with the same shape as ``state``.
    """
    xp = array_namespace(state)
    moved = xp.moveaxis(state, axis, -1)
    return xp.moveaxis(velocity * (moved @ generator), -1, axis)

axis_kernel_redistribute(flux, kernel, *, axis=-1)

Return flux @ kernel - flux along axis for a transferred flux.

Parameters:

Name	Type	Description	Default
flux	Any	Per-coordinate flux already moved by an ordinary transfer.	required
kernel	Any	Row-stochastic matrix over the kernel axis.	required
axis	int	Position of the kernel axis in flux.	-1

Returns:

Type	Description
Any	The redistribution to add to the target slice.

Source code in src/op_system/_axis_kernel.py

def axis_kernel_redistribute(
    flux: Any,  # ruff: ignore[any-type]
    kernel: Any,  # ruff: ignore[any-type]
    *,
    axis: int = -1,
) -> Any:  # ruff: ignore[any-type]
    """Return ``flux @ kernel - flux`` along ``axis`` for a transferred flux.

    Args:
        flux: Per-coordinate flux already moved by an ordinary transfer.
        kernel: Row-stochastic matrix over the kernel axis.
        axis: Position of the kernel axis in ``flux``.

    Returns:
        The redistribution to add to the target slice.
    """
    xp = array_namespace(flux)
    moved = xp.moveaxis(flux, axis, -1)
    return xp.moveaxis(moved @ kernel - moved, -1, axis)

compile_rhs(rhs, *, xp=None)

Compile a normalized RHS into a runnable evaluation function.

Always uses the vectorized eval path that operates on shaped buffers (one tensor expression per state template) for specs that declare axes. Specs without axes (genuinely scalar models) fall back to the scalar path. Raising :class:UnsupportedFeatureError if an axis-indexed spec cannot be vectorized, rather than silently falling back to the catastrophically slow scalar path.

The returned eval_fn is namespace-polymorphic: it infers its array namespace from the input y at call time (through array_api_compat.array_namespace), and returns arrays in that same namespace. Calling it with JAX arrays (or tracers) yields a JAX-native computation suitable for jax.jit / jax.vmap without any correctness wrapping.

Parameters:

Name	Type	Description	Default
rhs	NormalizedRhs	Normalized RHS produced by op_system.specs.normalize_rhs.	required
xp	object | None	Deprecated. Formerly the compile-time array backend namespace. Now ignored — the namespace is resolved per call from the input y. Will be removed in a future release.	None

Returns:

Type	Description
CompiledRhs	A CompiledRhs containing an eval_fn(t, y, **params) -> dydt.
CompiledRhs	For axis-indexed specs the returned object also carries
CompiledRhs	pytree_eval_fn and template_shapes. If the spec declares
CompiledRhs	axes but the vectorizer cannot build a plan an
CompiledRhs	UnsupportedFeatureError is raised (see bail reason in the detail
CompiledRhs	message) rather than silently degrading to the scalar path.

Source code in src/op_system/compile.py

def compile_rhs(rhs: NormalizedRhs, *, xp: object | None = None) -> CompiledRhs:
    """Compile a normalized RHS into a runnable evaluation function.

    Always uses the vectorized eval path that operates on shaped buffers
    (one tensor expression per state template) for specs that declare axes.
    Specs without axes (genuinely scalar models) fall back to the scalar path.
    Raising :class:`UnsupportedFeatureError` if an axis-indexed spec cannot be
    vectorized, rather than silently falling back to the catastrophically slow
    scalar path.

    The returned ``eval_fn`` is **namespace-polymorphic**: it infers its
    array namespace from the input ``y`` at call time
    (through ``array_api_compat.array_namespace``), and returns arrays in that same
    namespace. Calling it with JAX arrays (or tracers) yields a JAX-native
    computation suitable for ``jax.jit`` / ``jax.vmap`` without any
    correctness wrapping.

    Args:
        rhs: Normalized RHS produced by `op_system.specs.normalize_rhs`.
        xp: **Deprecated.** Formerly the compile-time array backend
            namespace. Now ignored — the namespace is resolved per call
            from the input ``y``. Will be removed in a future release.

    Returns:
        A `CompiledRhs` containing an `eval_fn(t, y, **params) -> dydt`.
        For axis-indexed specs the returned object also carries
        ``pytree_eval_fn`` and ``template_shapes``.  If the spec declares
        axes but the vectorizer cannot build a plan an
        ``UnsupportedFeatureError`` is raised (see bail reason in the detail
        message) rather than silently degrading to the scalar path.
    """
    _warn_on_deprecated_xp(xp)
    _validate_rhs_type(rhs)

    raw_history_requirements = (
        _history_requirements_from_ir(
            aliases_ir=rhs.aliases_ir,
            equations_ir=rhs.equations_ir,
        )
        if _reduce_ir_has_history(rhs)
        else ()
    )
    _validate_history_kinds(raw_history_requirements)

    vec, plan, eval_fn, pytree_eval_fn, template_shapes = _build_primary_eval_artifacts(
        rhs
    )

    # Apply both time-varying and synth-const wrappers to ``pytree_eval_fn``
    # BEFORE ``_build_history_artifacts`` consumes it.  history_eval_fn /
    # body_eval_fn capture the pytree_eval_fn reference at construction
    # time, so any later re-wrapping would not propagate into them; that
    # would leave runtime history-body calls missing time-varying param
    # slicing and synthesized constants (e.g. __op_system_mask__* one-hot
    # arrays for pinned transition selectors).
    synth_consts: Mapping[str, object] | None = None
    if eval_fn is not None:
        eval_fn, pytree_eval_fn = _wrap_time_varying_artifacts(
            rhs=rhs,
            eval_fn=eval_fn,
            pytree_eval_fn=pytree_eval_fn,
        )
        eval_fn, pytree_eval_fn, synth_consts = _apply_synth_const_wrappers(
            rhs=rhs,
            eval_fn=eval_fn,
            pytree_eval_fn=pytree_eval_fn,
        )

    history_requirements, history_eval_fn, body_eval_fn = _build_history_artifacts(
        rhs=rhs,
        plan=plan,
        pytree_eval_fn=pytree_eval_fn,
    )

    eval_fn = _resolve_eval_fn(
        rhs=rhs,
        eval_fn=eval_fn,
        history_requirements=history_requirements,
    )

    block_axes = analyze_block_axes(rhs)

    # ------------------------------------------------------------------
    # Block-stripped compile: produce a per-block-coord pytree_eval_fn by
    # stripping the first factorize axis from the RHS and re-running the
    # vectorizer.  Engines can jax.vmap this over the block axis instead
    # of baking literal axis indices that break under vmap.
    # ------------------------------------------------------------------
    block_pytree_eval_fn, block_template_shapes = _build_block_pytree_artifacts(
        rhs=rhs,
        vec=vec,
        pytree_eval_fn=pytree_eval_fn,
        block_axes=block_axes,
        synth_consts=synth_consts,
    )

    # Build per-block history / body eval fns when both the block compile and
    # history path succeeded.  These wrap ``block_pytree_eval_fn`` exactly as
    # ``history_eval_fn`` / ``body_eval_fn`` wrap ``pytree_eval_fn``.
    block_history_eval_fn, block_body_eval_fn = _build_block_history_artifacts(
        block_pytree_eval_fn=block_pytree_eval_fn,
        history_requirements=history_requirements,
    )

    return CompiledRhs(
        state_names=rhs.state_names,
        param_names=tuple(rhs.param_names),
        eval_fn=eval_fn,
        meta=rhs.meta,
        operators=_parse_operator_descriptors(rhs.meta),
        factorize_axes=_parse_factorize_axes(rhs.meta),
        block_axes=block_axes,
        pytree_eval_fn=pytree_eval_fn,
        template_shapes=template_shapes,
        block_pytree_eval_fn=block_pytree_eval_fn,
        block_template_shapes=block_template_shapes,
        history_requirements=history_requirements,
        history_eval_fn=history_eval_fn,
        body_eval_fn=body_eval_fn,
        block_history_eval_fn=block_history_eval_fn,
        block_body_eval_fn=block_body_eval_fn,
        reactions=_build_reaction_artifacts(rhs=rhs, plan=plan, vec=vec),
        _rhs=rhs,
    )

compile_spec(spec, *, xp=None, backend=DEFAULT_ARRAY_BACKEND)

Validate, normalize, and compile a RHS specification in one call.

This is the recommended public entrypoint for most users and adapters.

The compiled eval_fn is namespace-polymorphic: it infers its array namespace from the input y at call time (through array_api_compat.array_namespace), so one callable handles NumPy, JAX (concrete and traced), and any other Array-API backend natively. No compile-time backend selection is required.

Parameters:

Name	Type	Description	Default
spec	dict[str, object]	Raw RHS specification mapping (YAML/JSON friendly).	required
xp	object | None	Deprecated. Formerly the compile-time array backend namespace. Now ignored — see compile_rhs for details. Will be removed in a future release.	None
backend	Literal['numpy', 'jax']	Deprecated. Formerly selected the compile-time backend ("numpy" or "jax"). Now ignored. Will be removed in a future release.	DEFAULT_ARRAY_BACKEND

Returns:

Name	Type	Description
CompiledRhs	CompiledRhs	Runnable RHS callable container.

Source code in src/op_system/__init__.py

def compile_spec(  # ruff: ignore[non-empty-init-module]
    spec: dict[str, object],
    *,
    xp: object | None = None,
    backend: Literal["numpy", "jax"] = DEFAULT_ARRAY_BACKEND,
) -> CompiledRhs:
    """
    Validate, normalize, and compile a RHS specification in one call.

    This is the recommended public entrypoint for most users and adapters.

    The compiled ``eval_fn`` is **namespace-polymorphic**: it infers its
    array namespace from the input ``y`` at call time
    (through ``array_api_compat.array_namespace``), so one callable handles
    NumPy, JAX (concrete and traced), and any other Array-API backend
    natively. No compile-time backend selection is required.

    Args:
        spec: Raw RHS specification mapping (YAML/JSON friendly).
        xp: **Deprecated.** Formerly the compile-time array backend
            namespace. Now ignored — see ``compile_rhs`` for details.
            Will be removed in a future release.
        backend: **Deprecated.** Formerly selected the compile-time
            backend (``"numpy"`` or ``"jax"``). Now ignored. Will be
            removed in a future release.

    Returns:
        CompiledRhs: Runnable RHS callable container.
    """
    if xp is not None or backend != DEFAULT_ARRAY_BACKEND:
        warnings.warn(
            "compile_spec(xp=..., backend=...) is deprecated and ignored. "
            "The compiled eval_fn now infers its array namespace from the "
            "input `y` at call time; pass JAX arrays for a JAX-native call, "
            "NumPy arrays for a NumPy call, or Torch tensors for a Torch "
            "call. These kwargs will be removed in a future release.",
            DeprecationWarning,
            stacklevel=2,
        )

    rhs = normalize_rhs(spec)
    return compile_rhs(rhs)

jump_integral_generator(kernel, *, axis_type, direction='both', boundary='reflecting', quadrature_weights=None)

Build the conservative row-source generator for a jump kernel.

kernel[i, j] is the non-negative rate density from source coordinate i to target coordinate j. Its diagonal is ignored. up keeps only j > i; down keeps only j < i; both keeps every off-diagonal entry. For a continuous axis each target column j is multiplied by quadrature_weights[j] before the diagonal loss is set.

Parameters:

Name	Type	Description	Default
kernel	Any	Square source-by-target rate-density matrix.	required
axis_type	str	categorical, ordinal, or continuous.	required
direction	str	Permitted movement in declared coordinate order.	'both'
boundary	str	Boundary contract. Only reflecting is currently defined.	'reflecting'
quadrature_weights	Any | None	Positive target weights required for continuous axes and forbidden for categorical or ordinal axes.	None

Returns:

Type	Description
Any	A generator in the kernel's Array-API namespace. Rows sum to zero.

Raises:

Type	Description
ValueError	If static choices, shapes, or quadrature usage are invalid.

Notes

Call :func:validate_jump_integral_kernel with concrete parameter values at an orchestration boundary to check finiteness and signs. This builder deliberately keeps array values dynamic under transforms.

Source code in src/op_system/_jump_integral.py

def jump_integral_generator(
    kernel: Any,  # ruff: ignore[any-type]
    *,
    axis_type: str,
    direction: str = "both",
    boundary: str = "reflecting",
    quadrature_weights: Any | None = None,  # ruff: ignore[any-type]
) -> Any:  # ruff: ignore[any-type]
    """Build the conservative row-source generator for a jump kernel.

    ``kernel[i, j]`` is the non-negative rate density from source coordinate
    ``i`` to target coordinate ``j``. Its diagonal is ignored. ``up`` keeps
    only ``j > i``; ``down`` keeps only ``j < i``; ``both`` keeps every
    off-diagonal entry. For a continuous axis each target column ``j`` is
    multiplied by ``quadrature_weights[j]`` before the diagonal loss is set.

    Args:
        kernel: Square source-by-target rate-density matrix.
        axis_type: ``categorical``, ``ordinal``, or ``continuous``.
        direction: Permitted movement in declared coordinate order.
        boundary: Boundary contract. Only ``reflecting`` is currently defined.
        quadrature_weights: Positive target weights required for continuous
            axes and forbidden for categorical or ordinal axes.

    Returns:
        A generator in the kernel's Array-API namespace. Rows sum to zero.

    Raises:
        ValueError: If static choices, shapes, or quadrature usage are invalid.

    Notes:
        Call :func:`validate_jump_integral_kernel` with concrete parameter
        values at an orchestration boundary to check finiteness and signs.
        This builder deliberately keeps array values dynamic under transforms.
    """
    _validate_static_contract(
        axis_type=axis_type,
        direction=direction,
        boundary=boundary,
    )
    if kernel.ndim != 2 or kernel.shape[0] != kernel.shape[1]:
        msg = "jump_integral kernel must be square"
        raise ValueError(msg)

    xp = array_namespace(kernel)
    upper = xp.triu(xp.ones_like(kernel, dtype=xp.bool), k=1)
    lower = xp.tril(xp.ones_like(kernel, dtype=xp.bool), k=-1)
    if direction == "up":
        allowed = upper
    elif direction == "down":
        allowed = lower
    else:
        allowed = xp.logical_or(upper, lower)
    rates = xp.where(allowed, kernel, xp.zeros_like(kernel))

    if axis_type == "continuous":
        if quadrature_weights is None:
            msg = "continuous jump_integral axes require quadrature_weights"
            raise ValueError(msg)
        weights = xp.asarray(
            quadrature_weights,
            dtype=kernel.dtype,
            **_device_kwargs(kernel),
        )
        if weights.ndim != 1 or weights.shape[0] != kernel.shape[0]:
            msg = "quadrature_weights must have one entry per kernel target"
            raise ValueError(msg)
        rates *= xp.reshape(weights, (1, weights.shape[0]))
    elif quadrature_weights is not None:
        msg = "quadrature_weights are only valid for continuous jump_integral axes"
        raise ValueError(msg)

    row_sums = xp.sum(rates, axis=1)
    identity = xp.eye(
        kernel.shape[0],
        dtype=kernel.dtype,
        **_device_kwargs(kernel),
    )
    return rates - identity * xp.reshape(row_sums, (row_sums.shape[0], 1))

jump_integral_rhs(state, kernel, *, axis, axis_type, rate=1.0, direction='both', boundary='reflecting', quadrature_weights=None)

Return a conservative jump-integral derivative contribution.

Parameters:

Name	Type	Description	Default
state	Any	State array containing the jump axis.	required
kernel	Any	Source-by-target rate-density matrix.	required
axis	int	Position of the jump axis in state.	required
axis_type	str	categorical, ordinal, or continuous.	required
rate	Any	Scalar multiplier with inverse-time units after any continuous target quadrature has been applied.	1.0
direction	str	up, down, or both.	'both'
boundary	str	Boundary contract; currently only reflecting.	'reflecting'
quadrature_weights	Any | None	Target quadrature weights for continuous axes.	None

Returns:

Type	Description
Any	Derivative contribution with the same shape as state.

Source code in src/op_system/_jump_integral.py

def jump_integral_rhs(  # ruff: ignore[too-many-arguments]
    state: Any,  # ruff: ignore[any-type]
    kernel: Any,  # ruff: ignore[any-type]
    *,
    axis: int,
    axis_type: str,
    rate: Any = 1.0,  # ruff: ignore[any-type]
    direction: str = "both",
    boundary: str = "reflecting",
    quadrature_weights: Any | None = None,  # ruff: ignore[any-type]
) -> Any:  # ruff: ignore[any-type]
    """Return a conservative jump-integral derivative contribution.

    Args:
        state: State array containing the jump axis.
        kernel: Source-by-target rate-density matrix.
        axis: Position of the jump axis in ``state``.
        axis_type: ``categorical``, ``ordinal``, or ``continuous``.
        rate: Scalar multiplier with inverse-time units after any continuous
            target quadrature has been applied.
        direction: ``up``, ``down``, or ``both``.
        boundary: Boundary contract; currently only ``reflecting``.
        quadrature_weights: Target quadrature weights for continuous axes.

    Returns:
        Derivative contribution with the same shape as ``state``.
    """
    generator = jump_integral_generator(
        kernel,
        axis_type=axis_type,
        direction=direction,
        boundary=boundary,
        quadrature_weights=quadrature_weights,
    )
    return axis_kernel_generator_rhs(state, generator, axis=axis, velocity=rate)

normalize_expr_rhs(spec)

Normalize an expression-based RHS specification.

Parameters:

Name	Type	Description	Default
spec	Mapping[str, Any]	Raw RHS specification mapping.	required

Returns:

Type	Description
ExprRhs	Backend-facing normalized RHS representation.

Raises:

Type	Description
InvalidRhsSpecError	If validation fails.

Source code in src/op_system/_normalize.py

def normalize_expr_rhs(spec: Mapping[str, Any]) -> ExprRhs:  # ruff: ignore[complex-structure, too-many-locals, too-many-statements]
    """Normalize an expression-based RHS specification.

    Args:
        spec: Raw RHS specification mapping.

    Returns:
        Backend-facing normalized RHS representation.

    Raises:
        InvalidRhsSpecError: If validation fails.
    """
    state_raw = _ensure_str_list(spec.get("state"), name="state")
    if len(state_raw) != len(set(state_raw)):
        raise InvalidRhsSpecError(detail="state contains duplicate names")

    equations_map = spec.get("equations")
    if not isinstance(equations_map, dict):
        raise InvalidRhsSpecError(detail="equations must be a mapping of state->expr")

    equations_map = {_normalize_bracket_key(k): v for k, v in equations_map.items()}

    axes_meta = _normalize_axes(spec.get("axes"))
    meta_parts = _normalize_common_meta(
        spec,
        axis_names={"subgroup"} | {ax["name"] for ax in axes_meta},
        state_set=set(state_raw),
        operator_state_set=set(state_raw),
        axes=axes_meta,
    )

    meta: dict[str, Any] = {
        "axes": axes_meta,
        "state_axes": meta_parts[1],
        "kernels": meta_parts[2],
        "operators": meta_parts[3],
    }
    meta["op_system_may_have_history"] = _may_have_history(spec)
    for reserved_key in ("sources", "couplings", "constraints"):
        if reserved_key in spec:
            meta[reserved_key] = spec.get(reserved_key)

    state_expanded, state_template_map = _expand_state_templates(
        state_raw, axes=axes_meta
    )
    if len(state_expanded) != len(set(state_expanded)):
        raise InvalidRhsSpecError(detail="expanded state contains duplicates")

    # Pre-scan raw aliases + equations for shaped-parameter references before
    # alias/template expansion mangles bracketed bases into per-coord names.
    axis_lookup_dict: dict[str, list[str]] = build_axis_lookup(axes_meta)
    aliases_raw_map = meta_parts[0] or {}
    name_blocklist = (
        {parse_selector(s)[0] for s in state_raw}
        | set(state_expanded)
        | {parse_selector(_normalize_bracket_key(k))[0] for k in aliases_raw_map}
        | set(aliases_raw_map.keys())
    )
    raw_expressions: list[str] = [
        v for v in aliases_raw_map.values() if isinstance(v, str)
    ]
    raw_expressions.extend(v for v in equations_map.values() if isinstance(v, str))
    shaped_params = _scan_shaped_param_refs(
        raw_expressions,
        name_blocklist=name_blocklist,
        axis_lookup=axis_lookup_dict,
    )
    _reject_legacy_time_varying_field(spec)
    time_axis_name = _resolve_time_axis_name(spec)
    shaped_params, time_varying_full = _partition_time_varying_shaped(
        shaped_params,
        time_axis_name=time_axis_name,
        axis_lookup=axis_lookup_dict,
    )
    if time_varying_full:
        _strip_time_axis_in_mapping(
            equations_map,
            tv_full_axes=time_varying_full,
            time_axis_name=time_axis_name,
        )
        if isinstance(meta_parts[0], dict):
            _strip_time_axis_in_mapping(
                meta_parts[0],
                tv_full_axes=time_varying_full,
                time_axis_name=time_axis_name,
            )

    aliases_ir_map, aliases_ir_reduce_map, alias_template_map = (
        _build_aliases_ir_from_raw(
            meta_parts[0],
            axes=axes_meta,
            shaped_params=shaped_params,
            axis_lookup=axis_lookup_dict,
        )
    )
    template_map_all = {**state_template_map, **alias_template_map}

    chain_block = spec.get("chain")
    if chain_block:
        if not isinstance(chain_block, list):
            raise InvalidRhsSpecError(detail="chain must be a list if provided")
        _apply_expr_chains(
            chains=chain_block,
            state_expanded=state_expanded,
            equations_map=equations_map,
        )

    unknown_keys = [
        k
        for k in equations_map
        if k not in state_expanded and k not in template_map_all
    ]
    if unknown_keys:
        raise InvalidRhsSpecError(
            detail=f"unknown equation key(s): {sorted(unknown_keys)}"
        )

    equations_ir_built, equations_ir_reduce, all_syms = _build_equations_ir_from_raw(
        state_expanded=state_expanded,
        equations_map=equations_map,
        template_map=template_map_all,
        axes=axes_meta,
        shaped_params=shaped_params,
        axis_lookup=axis_lookup_dict,
        aliases_ir=aliases_ir_map,
    )
    # Collect free symbols from alias bodies (alias bodies may reference
    # params not appearing in any equation directly). Walk the *reduce*
    # map (Reduce nodes still folded) rather than the fully expanded
    # map: alias inlining is identical between the two, but the reduce
    # form is orders of magnitude smaller for continuum specs. Share a
    # single id-keyed memo across the per-cell entries so common
    # subtrees are visited at most once (issue #145).
    fs_memo: dict[int, frozenset[str]] = {}
    for alias_ir_val in aliases_ir_reduce_map.values():
        all_syms |= free_symbols(alias_ir_val, memo=fs_memo)

    _maybe_attach_initial_state(
        meta,
        spec.get("initial_state"),
        axes=axes_meta,
        template_map=template_map_all,
    )

    meta["shaped_params"] = tuple(sorted(shaped_params.items()))
    meta["time_axis"] = time_axis_name
    meta["time_varying_params"] = tuple(sorted(time_varying_full.items()))

    factorize_axes_raw = spec.get("factorize_axes")
    if factorize_axes_raw:
        known_axes = {ax["name"] for ax in axes_meta}
        meta["factorize_axes"] = [
            a for a in factorize_axes_raw if isinstance(a, str) and a in known_axes
        ]

    shaped_set = set(shaped_params)
    time_varying_set = set(time_varying_full)
    axis_name_set = set(axis_lookup_dict)
    template_base_set = {parse_selector(k)[0] for k in template_map_all}
    # Share an identity-keyed unparse memo between the equation and
    # alias rendering passes so subexpressions that recur across alias
    # bodies and equations (e.g. coord-pinned copies of the same
    # alias template body) are rendered only once (issue #145).
    unparse_memo: dict[tuple[int, int, bool], str] = {}
    equations_strings = _derive_equation_strings_lazy(
        equations_ir_built, _unparse_memo=unparse_memo
    )
    aliases_strings = _derive_alias_strings(
        aliases_ir_map, aliases_ir_map, _unparse_memo=unparse_memo
    )

    return ExprRhs(
        state_names=tuple(state_expanded),
        equations=equations_strings,
        aliases=aliases_strings,
        param_names=_sorted_unique(
            sym
            for sym in all_syms
            if sym not in set(state_expanded)
            and sym not in aliases_ir_map
            and sym not in shaped_set
            and sym not in time_varying_set
            and sym not in axis_name_set
            and sym not in template_base_set
            and sym not in _SHAPED_PARAM_BUILTIN_NAMES
        ),
        all_symbols=frozenset(all_syms | set(aliases_ir_map.keys())),
        meta=meta,
        state_templates=_build_state_templates(
            state_raw,
            axes=axes_meta,
            state_template_map=state_template_map,
            state_expanded=state_expanded,
        ),
        shaped_params=tuple(sorted(shaped_params.items())),
        time_varying_params=tuple(sorted(time_varying_full.items())),
        aliases_ir=aliases_ir_map,
        equations_ir=equations_ir_built,
        aliases_ir_reduce=aliases_ir_reduce_map,
        equations_ir_reduce=equations_ir_reduce,
        alias_templates=_build_alias_templates(
            aliases_raw_map,
            axes=axes_meta,
            alias_template_map=alias_template_map,
        ),
    )

normalize_rhs(spec)

Normalize a RHS specification dict into a backend-facing representation.

Parameters:

Name	Type	Description	Default
spec	Mapping[str, Any] | None	Raw RHS specification mapping.	required

Returns:

Type	Description
NormalizedRhs	Backend-facing normalized RHS representation.

Raises:

Type	Description
InvalidRhsSpecError	If validation fails.
UnsupportedFeatureError	If validation fails.

Source code in src/op_system/_normalize.py

def normalize_rhs(spec: Mapping[str, Any] | None) -> NormalizedRhs:
    """Normalize a RHS specification dict into a backend-facing representation.

    Args:
        spec: Raw RHS specification mapping.

    Returns:
        Backend-facing normalized RHS representation.

    Raises:
        InvalidRhsSpecError: If validation fails.
        UnsupportedFeatureError: If validation fails.
    """
    if spec is None:
        raise InvalidRhsSpecError(detail="rhs specification is required")

    kind = str(spec.get("kind", "expr")).strip().lower()

    if kind == "expr":
        return normalize_expr_rhs(spec)

    if kind == "transitions":
        return normalize_transitions_rhs(spec)

    raise UnsupportedFeatureError(
        feature=f"rhs.kind={kind}",
        detail="Only 'expr' and 'transitions' are supported in v1.",
    )

normalize_transitions_rhs(spec)

Normalize a transition-based RHS specification (diagram/hazard semantics).

Returns:

Type	Description
TransitionsRhs	Backend-facing normalized RHS representation for the transitions kind.

Raises:

Type	Description
InvalidRhsSpecError	If validation fails.

Source code in src/op_system/_normalize.py

def normalize_transitions_rhs(  # ruff: ignore[complex-structure, too-many-branches, too-many-locals, too-many-statements]
    spec: Mapping[str, Any],
) -> TransitionsRhs:
    """Normalize a transition-based RHS specification (diagram/hazard semantics).

    Returns:
        Backend-facing normalized RHS representation for the transitions kind.

    Raises:
        InvalidRhsSpecError: If validation fails.
    """
    state_raw = _ensure_str_list(spec.get("state"), name="state")
    if len(state_raw) != len(set(state_raw)):
        raise InvalidRhsSpecError(detail="state contains duplicate names")

    transitions_raw = spec.get("transitions")
    if transitions_raw is None:
        transitions_raw = []
    elif isinstance(transitions_raw, list):
        # Copy each entry: time-axis stripping rewrites rates in place.
        transitions_raw = [
            dict(tr) if isinstance(tr, _MappingABC) else tr for tr in transitions_raw
        ]
    else:
        raise InvalidRhsSpecError(detail="transitions must be a list")

    axes_meta = _normalize_axes(spec.get("axes"))

    meta_parts = _normalize_common_meta(
        spec,
        axis_names={"subgroup"} | {ax["name"] for ax in axes_meta},
        state_set=None,
        operator_state_set=set(state_raw),
        axes=axes_meta,
    )

    meta: dict[str, Any] = {
        "transitions": transitions_raw,
        "axes": axes_meta,
        "kernels": meta_parts[2],
        "operators": meta_parts[3],
    }
    meta["op_system_may_have_history"] = _may_have_history(spec)
    meta.update({
        k: spec[k] for k in ("sources", "couplings", "constraints") if k in spec
    })

    chain_block = spec.get("chain")
    if chain_block:
        if not isinstance(chain_block, list):
            raise InvalidRhsSpecError(detail="chain must be a list if provided")
        _apply_transition_chains(
            chains=chain_block,
            state_raw=state_raw,
            transitions_raw=transitions_raw,
            state_set=set(state_raw),
        )

    state_expanded, state_template_map = _expand_state_templates(
        state_raw, axes=axes_meta
    )
    if len(state_expanded) != len(set(state_expanded)):
        raise InvalidRhsSpecError(detail="expanded state contains duplicates")

    # Pre-scan raw aliases + transition rates for shaped-parameter references.
    axis_lookup_dict: dict[str, list[str]] = build_axis_lookup(axes_meta)
    aliases_raw_map = meta_parts[0] or {}
    name_blocklist = (
        {parse_selector(s)[0] for s in state_raw}
        | set(state_expanded)
        | {parse_selector(_normalize_bracket_key(k))[0] for k in aliases_raw_map}
        | set(aliases_raw_map.keys())
    )
    raw_expressions: list[str] = [
        v for v in aliases_raw_map.values() if isinstance(v, str)
    ]
    for tr in transitions_raw:
        if isinstance(tr, _MappingABC):
            r = tr.get("rate")
            if isinstance(r, str):
                raw_expressions.append(r)
            n = tr.get("name")
            if isinstance(n, str):
                raw_expressions.append(n)
    shaped_params = _scan_shaped_param_refs(
        raw_expressions,
        name_blocklist=name_blocklist,
        axis_lookup=axis_lookup_dict,
    )
    _reject_legacy_time_varying_field(spec)
    time_axis_name = _resolve_time_axis_name(spec)
    shaped_params, time_varying_full = _partition_time_varying_shaped(
        shaped_params,
        time_axis_name=time_axis_name,
        axis_lookup=axis_lookup_dict,
    )
    if time_varying_full:
        if isinstance(meta_parts[0], dict):
            _strip_time_axis_in_mapping(
                meta_parts[0],
                tv_full_axes=time_varying_full,
                time_axis_name=time_axis_name,
            )
        for tr in transitions_raw:
            if not isinstance(tr, _MappingABC):
                continue
            r = tr.get("rate")
            if isinstance(r, str):
                tr["rate"] = _strip_time_axis_in_expr(  # type: ignore[index]
                    r,
                    tv_full_axes=time_varying_full,
                    time_axis_name=time_axis_name,
                )
            n = tr.get("name")
            if isinstance(n, str):
                tr["name"] = _strip_time_axis_in_expr(  # type: ignore[index]
                    n,
                    tv_full_axes=time_varying_full,
                    time_axis_name=time_axis_name,
                )

    aliases_ir_map, aliases_ir_reduce_map, alias_template_map = (
        _build_aliases_ir_from_raw(
            aliases_raw_map,
            axes=axes_meta,
            shaped_params=shaped_params,
            axis_lookup=axis_lookup_dict,
        )
    )
    template_map_all = {**state_template_map, **alias_template_map}

    state_set = set(state_expanded)
    # Collect alias symbols from IR. Walk the Reduce-folded map (same
    # inlined symbol set as the full-expansion map but vastly smaller
    # on continuum specs) and share an id-keyed memo across per-cell
    # entries that share subtrees by identity (issue #145).
    all_syms: set[str] = set()
    fs_memo: dict[int, frozenset[str]] = {}
    for alias_expr in aliases_ir_reduce_map.values():
        all_syms |= free_symbols(alias_expr, memo=fs_memo)

    _apply_coord_shifts(
        transitions_raw=transitions_raw,
        state_expanded=state_expanded,
        axes=axes_meta,
        state_template_map=state_template_map,
    )

    if not transitions_raw:
        raise InvalidRhsSpecError(
            detail="transitions must be non-empty after applying chain expansion"
        )

    # Build equations and collect expanded transitions in one IR-native pass
    pinned_mask_names, pinned_mask_values = _discover_pinned_token_masks(
        transitions_raw, axis_lookup=axis_lookup_dict
    )
    if pinned_mask_values:
        # Register one-hot masks as shaped params so the vectorizer's
        # extra-param-buffers plumbing assembles them at eval time; stash
        # the actual values under ``meta`` so ``compile_rhs`` can inject
        # them into the eval_fn's ``params`` automatically.
        for (axis, _coord), mask_name in pinned_mask_names.items():
            shaped_params[mask_name] = (axis,)
        meta["op_system_synth_constants"] = dict(pinned_mask_values)

    equations_ir_pre_inline, equations_ir_reduce, transitions_expanded, rate_syms = (
        _build_transition_equations_ir(
            transitions_raw,
            state_set=state_set,
            state_expanded=state_expanded,
            axes=axes_meta,
            axis_lookup=axis_lookup_dict,
            template_map=template_map_all,
            shaped_params=shaped_params,
            mask_names=pinned_mask_names,
            time_axis_name=time_axis_name,
            alias_bases={
                parse_selector(_normalize_bracket_key(k))[0] for k in aliases_raw_map
            },
        )
    )
    all_syms |= rate_syms

    reactions_ir = build_reaction_artifacts_ir(
        transitions_raw,
        axes=axes_meta,
        axis_lookup=axis_lookup_dict,
        shaped_params=shaped_params,
        time_axis_name=time_axis_name,
        aliases_raw=aliases_raw_map,
        state_axes={
            base: tuple(tok.axis for tok in tokens)
            for state_s in state_raw
            for base, tokens in (parse_selector(state_s),)
        },
    )

    _maybe_attach_initial_state(
        meta,
        spec.get("initial_state"),
        axes=axes_meta,
        template_map=template_map_all,
    )

    meta["shaped_params"] = tuple(sorted(shaped_params.items()))
    meta["time_axis"] = time_axis_name
    meta["time_varying_params"] = tuple(sorted(time_varying_full.items()))

    factorize_axes_raw = spec.get("factorize_axes")
    if factorize_axes_raw:
        known_axes = {ax["name"] for ax in axes_meta}
        meta["factorize_axes"] = [
            a for a in factorize_axes_raw if isinstance(a, str) and a in known_axes
        ]

    shaped_set = set(shaped_params)
    time_varying_set = set(time_varying_full)
    axis_name_set = set(axis_lookup_dict)
    template_base_set = {parse_selector(k)[0] for k in template_map_all}
    # Share an identity-keyed unparse memo between the equation and
    # alias rendering passes so subexpressions that recur across alias
    # bodies and equations are rendered only once (issue #145).
    unparse_memo_final: dict[tuple[int, int, bool], str] = {}
    eqs_tuple = _derive_equation_strings(
        equations_ir_pre_inline, _unparse_memo=unparse_memo_final
    )
    # Inline aliases directly into the per-cell IR we already built in
    # ``_build_transition_equations_ir`` rather than re-parsing the
    # round-tripped equation strings. Avoids 73k x ``parse_expr_to_ir``
    # plus a redundant IR rebuild on large continuum specs (issue #145).
    # ``aliases_ir_map`` is already fully alias-inlined inside
    # ``_build_aliases_ir_from_raw`` so cycle detection is redundant here.
    #
    # The dominant cost on large specs is the per-cell ``inline_aliases``
    # call. Because synthesized transitions install the SAME ``synth_to`` /
    # ``synth_neg`` IR object into every cell of a template, the *terms*
    # of the per-cell sum are shared across many cells even though the
    # outer ``Apply(op="+", ...)`` wrapper is unique per cell. Inlining
    # term-by-term with a shared ``result_memo`` keyed on ``id(term)``
    # collapses the alias-substitution work from O(n_state) to
    # O(n_unique_terms) (issue #145).
    # Both memos key ``id()`` of terms owned by ``equations_ir_pre_inline``,
    # which outlives them -- the invariant that makes an identity key sound
    # (issue #200). Passing a freshly-built expression to either would
    # silently return another expression's inlined result.
    alias_inline_memo = _InlineMemo()
    alias_inline_result_memo: dict[int, Expr] = {}

    def _inline_one(expr: Expr) -> Expr:
        return inline_aliases(
            expr,
            aliases_ir_map,
            memo=alias_inline_memo,
            skip_cycle_check=True,
            result_memo=alias_inline_result_memo,
        )

    # Dedup the post-inline outer Apply by tuple of arg ids so cells
    # whose inlined-term tuple is identical share one wrapper, letting
    # downstream identity-keyed memos (e.g. unparse_ir) cache the
    # rendered string once per unique equation. Issue #145.
    apply_plus_dedup: dict[tuple[int, ...], Expr] = {}
    # The upstream ``sum_dedup_full`` collapses ``equations_ir_pre_inline``
    # to a handful of unique outer Apply objects shared across many
    # cells (e.g. 7 unique exprs across 72,828 cells on the COVID19_USA
    # continuum spec). Cache the per-expr inlined result by ``id(expr)``
    # so we skip the per-term ``_inline_one`` calls entirely on cache
    # hits -- this collapses ~1.9M memo-hit calls to ~7 real inlines
    # plus 72k dict lookups (issue #145).
    outer_inline_memo: dict[int, Expr] = {}

    def _inline_outer(expr: Expr) -> Expr:
        """Run alias-inlining for one outer equation expression.

        Returns:
            The inlined expression (possibly the original ``expr`` when no
            inlining was needed).
        """
        if isinstance(expr, Apply) and expr.op == "+":
            new_args = tuple(_inline_one(a) for a in expr.args)
            if all(n is o for n, o in zip(new_args, expr.args, strict=True)):
                return expr
            dedup_key = tuple(id(a) for a in new_args)
            cached_apply = apply_plus_dedup.get(dedup_key)
            if cached_apply is None:
                cached_apply = Apply(op="+", args=new_args)
                apply_plus_dedup[dedup_key] = cached_apply
            return cached_apply
        return _inline_one(expr)

    equations_ir_built_list: list[Expr | None] = []
    for expr in equations_ir_pre_inline:
        if expr is None:
            equations_ir_built_list.append(None)
            continue
        if not aliases_ir_map:
            equations_ir_built_list.append(expr)
            continue
        cached_outer = outer_inline_memo.get(id(expr))
        if cached_outer is not None:
            equations_ir_built_list.append(cached_outer)
            continue
        try:
            result_expr = _inline_outer(expr)
        except (ValueError, RecursionError):
            result_expr = expr
        outer_inline_memo[id(expr)] = result_expr
        equations_ir_built_list.append(result_expr)
    equations_ir_built = tuple(equations_ir_built_list)
    return TransitionsRhs(
        reactions_ir=reactions_ir,
        state_names=tuple(state_expanded),
        equations=eqs_tuple,
        aliases=_derive_alias_strings(
            aliases_ir_map, aliases_ir_map, _unparse_memo=unparse_memo_final
        ),
        param_names=_sorted_unique(
            sym
            for sym in all_syms
            if sym not in state_set
            and sym not in aliases_ir_map
            and sym not in shaped_set
            and sym not in time_varying_set
            and sym not in axis_name_set
            and sym not in template_base_set
            and sym not in _SHAPED_PARAM_BUILTIN_NAMES
        ),
        all_symbols=frozenset(all_syms | set(aliases_ir_map.keys())),
        meta={**meta, "transitions": transitions_expanded},
        state_templates=_build_state_templates(
            state_raw,
            axes=axes_meta,
            state_template_map=state_template_map,
            state_expanded=state_expanded,
        ),
        shaped_params=tuple(sorted(shaped_params.items())),
        time_varying_params=tuple(sorted(time_varying_full.items())),
        aliases_ir=aliases_ir_map,
        equations_ir=equations_ir_built,
        aliases_ir_reduce=aliases_ir_reduce_map,
        equations_ir_reduce=equations_ir_reduce,
        alias_templates=_build_alias_templates(
            aliases_raw_map,
            axes=axes_meta,
            alias_template_map=alias_template_map,
        ),
    )

validate_axis_kernel_matrix(matrix, *, form, tolerance=1e-09)

Return problems with a numeric kernel matrix for form (empty if valid).

Parameters:

Name	Type	Description	Default
matrix	Any	Candidate square matrix.	required
form	str	"generator" or "stochastic".	required
tolerance	float	Absolute tolerance for sign and row-sum checks.	1e-09

Returns:

Type	Description
list[str]	Human-readable problems; an empty list means the matrix is valid.

Raises:

Type	Description
ValueError	If form is not a supported axis-kernel form.

Source code in src/op_system/_axis_kernel.py

def validate_axis_kernel_matrix(
    matrix: Any,  # ruff: ignore[any-type]
    *,
    form: str,
    tolerance: float = 1e-9,
) -> list[str]:
    """Return problems with a numeric kernel matrix for ``form`` (empty if valid).

    Args:
        matrix: Candidate square matrix.
        form: ``"generator"`` or ``"stochastic"``.
        tolerance: Absolute tolerance for sign and row-sum checks.

    Returns:
        Human-readable problems; an empty list means the matrix is valid.

    Raises:
        ValueError: If ``form`` is not a supported axis-kernel form.
    """
    if form not in AXIS_KERNEL_FORMS:
        msg = f"unknown axis_kernel form {form!r}"
        raise ValueError(msg)
    xp = array_namespace(matrix)
    problems: list[str] = []
    if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]:
        return ["matrix must be square"]
    n = matrix.shape[0]
    off_diagonal = matrix * (1.0 - xp.eye(n, dtype=matrix.dtype))
    if bool(xp.any(off_diagonal < -tolerance)):
        problems.append("off-diagonal entries must be nonnegative")
    row_sums = xp.sum(matrix, axis=1)
    if form == "generator":
        if bool(xp.any(xp.abs(row_sums) > tolerance)):
            problems.append("generator rows must sum to zero")
    else:
        if bool(xp.any(xp.linalg.diagonal(matrix) < -tolerance)):
            problems.append("stochastic diagonal entries must be nonnegative")
        if bool(xp.any(xp.abs(row_sums - 1.0) > tolerance)):
            problems.append("stochastic rows must sum to one")
    return problems

validate_jump_integral_kernel(kernel, *, axis_type, direction='both', boundary='reflecting', quadrature_weights=None, tolerance=1e-09)

Return problems with concrete jump-kernel values (empty if valid).

This eager validation belongs at parameter resolution or another host orchestration boundary, not inside a traced numerical solve.

Parameters:

Name	Type	Description	Default
kernel	Any	Candidate source-by-target rate-density matrix.	required
axis_type	str	categorical, ordinal, or continuous.	required
direction	str	up, down, or both.	'both'
boundary	str	Boundary contract; currently only reflecting.	'reflecting'
quadrature_weights	Any | None	Target quadrature weights for continuous axes.	None
tolerance	float	Absolute tolerance for sign and diagonal checks.	1e-09

Returns:

Type	Description
list[str]	Human-readable problems; an empty list means the values conform.

Raises:

Type	Description
ValueError	If a static contract choice is unsupported.

Source code in src/op_system/_jump_integral.py

def validate_jump_integral_kernel(  # ruff: ignore[complex-structure, too-many-arguments, too-many-branches]
    kernel: Any,  # ruff: ignore[any-type]
    *,
    axis_type: str,
    direction: str = "both",
    boundary: str = "reflecting",
    quadrature_weights: Any | None = None,  # ruff: ignore[any-type]
    tolerance: float = 1e-9,
) -> list[str]:
    """Return problems with concrete jump-kernel values (empty if valid).

    This eager validation belongs at parameter resolution or another host
    orchestration boundary, not inside a traced numerical solve.

    Args:
        kernel: Candidate source-by-target rate-density matrix.
        axis_type: ``categorical``, ``ordinal``, or ``continuous``.
        direction: ``up``, ``down``, or ``both``.
        boundary: Boundary contract; currently only ``reflecting``.
        quadrature_weights: Target quadrature weights for continuous axes.
        tolerance: Absolute tolerance for sign and diagonal checks.

    Returns:
        Human-readable problems; an empty list means the values conform.

    Raises:
        ValueError: If a static contract choice is unsupported.
    """
    try:
        _validate_static_contract(
            axis_type=axis_type,
            direction=direction,
            boundary=boundary,
        )
    except ValueError as exc:
        raise ValueError(str(exc)) from exc
    if kernel.ndim != 2 or kernel.shape[0] != kernel.shape[1]:
        return ["kernel must be square"]
    xp = array_namespace(kernel)
    problems: list[str] = []
    if bool(xp.any(xp.logical_not(xp.isfinite(kernel)))):
        problems.append("kernel entries must be finite")
    if bool(xp.any(kernel < -tolerance)):
        problems.append("kernel entries must be nonnegative")
    if bool(xp.any(xp.abs(xp.linalg.diagonal(kernel)) > tolerance)):
        problems.append("kernel diagonal entries must be zero")

    if axis_type == "continuous":
        if quadrature_weights is None:
            problems.append("continuous axes require quadrature_weights")
        else:
            weights = xp.asarray(
                quadrature_weights,
                dtype=kernel.dtype,
                **_device_kwargs(kernel),
            )
            if weights.ndim != 1 or weights.shape[0] != kernel.shape[0]:
                problems.append("quadrature_weights must match kernel size")
            else:
                if bool(xp.any(xp.logical_not(xp.isfinite(weights)))):
                    problems.append("quadrature_weights must be finite")
                if bool(xp.any(weights <= 0.0)):
                    problems.append("quadrature_weights must be positive")
    elif quadrature_weights is not None:
        problems.append("quadrature_weights require a continuous axis")
    return problems

validate_spec(spec)

Validate an op_system specification without raising.

Parameters:

Name	Type	Description	Default
spec	Mapping[str, Any]	A raw RHS specification mapping.	required

Returns:

Type	Description
ValidationReport	A report with per-stage status, error messages, compile-cost drivers,
ValidationReport	consumed parameters, and expression-shape groups for templates whose
ValidationReport	cells differ.

Source code in src/op_system/_validate.py

def validate_spec(spec: Mapping[str, Any]) -> ValidationReport:
    """Validate an op_system specification without raising.

    Args:
        spec: A raw RHS specification mapping.

    Returns:
        A report with per-stage status, error messages, compile-cost drivers,
        consumed parameters, and expression-shape groups for templates whose
        cells differ.
    """
    from op_system import compile_rhs, normalize_rhs  # ruff: ignore[import-outside-top-level]

    stages = dict.fromkeys(_STAGES, "skipped")
    try:
        rhs = normalize_rhs(spec)
    except Exception as exc:  # ruff: ignore[blind-except] - reported, not raised
        stages["normalize"] = "failed"
        return ValidationReport(
            stages=stages, errors=[str(exc)], cost=_cost(spec, None)
        )
    stages["normalize"] = "passed"
    report_kwargs: dict[str, Any] = {
        "cost": _cost(spec, rhs),
        "parameters": _parameters(rhs),
    }
    try:
        compile_rhs(rhs)
    except UnsupportedFeatureError as exc:
        vectorize = "vectorized eval path" in str(exc)
        stages["compile"] = "passed" if vectorize else "failed"
        stages["vectorize"] = "failed" if vectorize else "skipped"
        if vectorize:
            report_kwargs["shape_groups"] = _shape_groups(rhs)
        return ValidationReport(stages=stages, errors=[str(exc)], **report_kwargs)
    except Exception as exc:  # ruff: ignore[blind-except]
        stages["compile"] = "failed"
        return ValidationReport(stages=stages, errors=[str(exc)], **report_kwargs)
    stages["compile"] = "passed"
    stages["vectorize"] = "passed"
    return ValidationReport(stages=stages, **report_kwargs)