Specs

specs

op_system.specs.

Public façade for RHS specification models and normalization.

The implementation lives in op_system._normalize. This module re-exports the public types and functions under the stable op_system.specs namespace for backward compatibility.

Supported RHS kinds

1) kind: "expr" - explicit d(state)/dt equations per state variable. 2) kind: "transitions" - diagram-style per-capita hazard transitions.

See op_system._normalize for the full implementation.

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)

PinnedToken(axis, coord) dataclass

Axis token pinned to a specific coordinate.

StateTemplate(base, axes, shape, expanded_names, coord_assignments, offset) dataclass

Structural record for a state template prior to scalar expansion.

A spec like state: ["S[age, vax]"] over axes age (size 4) and vax (size 2) produces a single StateTemplate with base="S", axes=("age", "vax"), shape=(4, 2), and expanded_names listing the eight scalar state names in cartesian-product order (age outer, vax inner — matching itertools.product and the order used to expand the flat state vector).

Non-templated entries (e.g. a bare "D" in state) are also reported as StateTemplate records with axes=(), shape=(), and a single expanded_names = ("D",) so consumers can iterate templates uniformly.

Attributes:

Name	Type	Description
base	str	Compartment name without selector brackets (e.g. "S").
axes	tuple[str, ...]	Wildcard axes in declaration order. Empty for scalar templates.
shape	tuple[int, ...]	Per-axis sizes in axes order. Empty for scalar templates.
expanded_names	tuple[str, ...]	Flat scalar state names, in cartesian-product order over axes (consistent with state_names).
coord_assignments	tuple[Mapping[str, str], ...]	For each entry in expanded_names, the axis -> coord mapping. Empty dict for scalar templates.
offset	int	Index of expanded_names[0] within the parent NormalizedRhs.state_names tuple (i.e. where this template's slice starts in the flat state vector).

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=()) 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.

WildcardToken(axis) dataclass

Axis token that expands over all coordinates of that axis.

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:  # noqa: C901, PLR0914, PLR0915
    """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],
    }
    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(
        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(  # noqa: C901, PLR0912, PLR0914, PLR0915
    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):
        transitions_raw = list(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.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

    _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).
    alias_inline_memo: dict[int, frozenset[str]] = {}
    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(
        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,
        ),
    )