Code generation

OSQP-codegen-style generation of low-level solvers for verified first-order methods, plus a tiny Python wrapper. An LLM coding agent translates each step of the algorithm into the target language and iterates against a Python reference oracle until the generated code is numerically verified.

The code generation guide is the narrative reference for the symbols below, and the tutorials work through five complete examples.

generate_code

generate_code(verifier, next_state=None, out_dir: str | Path = None, config: CodegenConfig | None = None, driver: AgentDriver | None = None, state=None, params=None, name: str = 'fom', code_optimization: bool = False, profile: bool = False, schedules: dict | None = None, iterates=None, param_values: dict | None = None, roll: bool = True) -> CodegenResult

Generate a low-level solver for the FOM recorded in verifier.

Three ways in, selected by what you pass:

• iterates= with no next_state generates the whole recorded run, straight from the loop you wrote to verify with. This accepts arbitrary per-iteration coefficients and step types, and anchored methods. It rolls the run back up when repeating one iteration is verified to reproduce it, so a stationary method costs nothing here. See extract_unrolled_ir.
• next_state= generates a rolled solver from one stationary iteration, repeated config.termination.max_iters times.
• a build function plus schedules= generates a rolled solver whose declared coefficients are tabulated and read at the loop counter. See extract_scheduled_ir.

Parameters:

Name	Type	Description	Default
verifier		Forwarded to fomtools.codegen.extract.extract_ir. For a method with an iteration-varying coefficient, pass a build function as verifier and its schedules instead: the function records one iteration as a function of the schedule scalars and returns a Recording. See extract_scheduled_ir.	required
next_state		Forwarded to fomtools.codegen.extract.extract_ir. For a method with an iteration-varying coefficient, pass a build function as verifier and its schedules instead: the function records one iteration as a function of the schedule scalars and returns a Recording. See extract_scheduled_ir.	required
state		Forwarded to fomtools.codegen.extract.extract_ir. For a method with an iteration-varying coefficient, pass a build function as verifier and its schedules instead: the function records one iteration as a function of the schedule scalars and returns a Recording. See extract_scheduled_ir.	required
params		Forwarded to fomtools.codegen.extract.extract_ir. For a method with an iteration-varying coefficient, pass a build function as verifier and its schedules instead: the function records one iteration as a function of the schedule scalars and returns a Recording. See extract_scheduled_ir.	required
out_dir	str | Path	Directory to emit the solver into.	None
config	CodegenConfig | None	Code-generation configuration.	None
driver	AgentDriver | None	The coding-agent driver. If None, uses fomtools.codegen.agent.get_default_driver.	None
name	str	FOM name used for file naming.	'fom'
code_optimization	bool	Run the code optimization phase (iterative profile-guided rewriting) after a successful, verified translation.	False
profile	bool	Collect before/after timing + a profiling breakdown (needs a C compiler). Always skipped silently when no compiler is available.	False
schedules	dict | None	Iteration-varying scalars, mapping name to a callable of the 0-based iteration index, a sequence, or a Schedule. Selects the scheduled extraction path.	None
iterates		The successive recorded iterate expressions, or an IterateTrajectory. Passed without next_state this selects the unrolled path and is the run to generate. Passed with next_state it instead checks the rolled IR against the recorded tape before anything is emitted, failing if the recording is not stationary — strongly recommended there, since without it nothing detects a method whose coefficients vary with k.	None
param_values	dict | None	Parameter values, keyed by parameter leaf, used for whichever check the chosen path runs. Optional on the unrolled path, which draws its own.	None
roll	bool	Unrolled path only: emit the compact rolled IR when repeating one iteration is verified to reproduce the recording. Pass False to unroll regardless.	True

Source code in fomtools/codegen/generate.py

def generate_code(
    verifier,
    next_state=None,
    out_dir: str | Path = None,
    config: CodegenConfig | None = None,
    driver: AgentDriver | None = None,
    state=None,
    params=None,
    name: str = "fom",
    code_optimization: bool = False,
    profile: bool = False,
    schedules: dict | None = None,
    iterates=None,
    param_values: dict | None = None,
    roll: bool = True,
) -> CodegenResult:
    """Generate a low-level solver for the FOM recorded in ``verifier``.

    Three ways in, selected by what you pass:

    * ``iterates=`` with no ``next_state`` generates the **whole recorded run**,
      straight from the loop you wrote to verify with. This accepts arbitrary
      per-iteration coefficients and step types, and anchored methods. It rolls
      the run back up when repeating one iteration is verified to reproduce it,
      so a stationary method costs nothing here. See
      `extract_unrolled_ir`.
    * ``next_state=`` generates a **rolled** solver from one stationary
      iteration, repeated ``config.termination.max_iters`` times.
    * a *build function* plus ``schedules=`` generates a rolled solver whose
      declared coefficients are tabulated and read at the loop counter. See
      `extract_scheduled_ir`.

    Parameters
    ----------
    verifier, next_state, state, params:
        Forwarded to `fomtools.codegen.extract.extract_ir`.

        For a method with an iteration-varying coefficient, pass a *build
        function* as ``verifier`` and its ``schedules`` instead: the function
        records one iteration as a function of the schedule scalars and returns a
        `Recording`. See
        `extract_scheduled_ir`.
    out_dir:
        Directory to emit the solver into.
    config:
        Code-generation configuration.
    driver:
        The coding-agent driver. If ``None``, uses
        `fomtools.codegen.agent.get_default_driver`.
    name:
        FOM name used for file naming.
    code_optimization:
        Run the code optimization phase (iterative profile-guided rewriting)
        after a successful, verified translation.
    profile:
        Collect before/after timing + a profiling breakdown (needs a C
        compiler). Always skipped silently when no compiler is available.
    schedules:
        Iteration-varying scalars, mapping name to a callable of the 0-based
        iteration index, a sequence, or a
        `Schedule`. Selects the scheduled
        extraction path.
    iterates:
        The successive recorded iterate expressions, or an
        `IterateTrajectory`.

        Passed *without* ``next_state`` this selects the unrolled path and is the
        run to generate. Passed *with* ``next_state`` it instead checks the rolled
        IR against the recorded tape before anything is emitted, failing if the
        recording is not stationary — strongly recommended there, since without it
        nothing detects a method whose coefficients vary with ``k``.
    param_values:
        Parameter values, keyed by parameter *leaf*, used for whichever check the
        chosen path runs. Optional on the unrolled path, which draws its own.
    roll:
        Unrolled path only: emit the compact rolled IR when repeating one
        iteration is verified to reproduce the recording. Pass ``False`` to unroll
        regardless.
    """
    config = config or CodegenConfig()
    out_dir = Path(out_dir)

    ir = extract_fom_ir(
        verifier,
        next_state,
        config=config,
        state=state,
        params=params,
        name=name,
        schedules=schedules,
        iterates=iterates,
        param_values=param_values,
        roll=roll,
    )

    target = scaffold(ir, out_dir, config)

    if driver is None:
        from fomtools.codegen.agent import get_default_driver

        driver = get_default_driver(config)
    elif getattr(driver, "config", False) is None:
        # A driver built without a config (e.g. ClaudeAgentDriver()) takes this
        # run's budgets instead of its constructor default.
        driver.config = config

    from fomtools.codegen.agent.prompts import translation_instructions

    # The agent may edit only the solver body; restore all deterministic files
    # (Makefile, headers, numeric data, factorized solve, sparse matvecs) before
    # each verify so it cannot break the build infrastructure or the data.
    def restore_deterministic_files():
        target.scaffold(ir, out_dir)

    # In the code optimization phase the agent may also rewrite the linear solves
    # and the matrix-vector products (see run_code_optimization).
    editable = tuple(target.optimization_editable_paths())

    # --- Translation phase ----------------------------------------------------
    translation_task = AgentTask(
        workdir=out_dir,
        solver_rel_path=target.solver_source_path(),
        verify_cmd=target.verify_command(),
        instructions=translation_instructions(ir, target),
        phase="translation",
        restore_hook=restore_deterministic_files,
    )
    # `make verify` also holds the translated solver to the kernel rule that has
    # no baseline (it reads no matrix data: every product goes through a checked
    # fom_apply_*); the check below applies it again and decides.
    target.write_kernel_check_config(out_dir, ir, None)
    try:
        translate = driver.run(translation_task)
    finally:
        target.remove_kernel_check_config(out_dir)
    if translate.success:
        solver = (out_dir / target.solver_source_path()).read_text()
        issues = target.translation_checks(solver, ir)
        if issues:
            notes = "\n".join(f"[translation rejected: {v}]" for v in issues)
            translate = replace(translate, success=False, log=translate.log + "\n" + notes)

    timing = TimingReport() if profile else None
    have_cc = find_compiler() is not None

    # Baseline timing once a correct solver exists.
    if translate.success and profile and have_cc:
        restore_deterministic_files()  # ensure a clean build before measuring
        sample = representative_params(ir)
        try:
            timing.python_us = time_python_oracle(ir, sample)
        except Exception:
            timing.python_us = None
        timing.c_initial_us = _bench_mean_us(out_dir)
        timing.c_final_us = timing.c_initial_us  # updated after code optimization
        try:
            timing.profile = run_profile(out_dir)
        except Exception:
            timing.profile = []

    # --- Code optimization phase ---------------------------------------------
    optimize = None
    optimization_diff = None
    optimizations_md = None
    kernel_checks = None
    rule_friction: list[dict] = []
    if code_optimization and translate.success:
        opt = run_code_optimization(
            ir,
            target,
            out_dir,
            driver,
            config,
            profile=profile,
            timing=timing,
            translated_verify_output=translate.log,
        )
        optimize = opt.optimize
        optimization_diff = opt.diff
        optimizations_md = opt.optimizations_md
        kernel_checks = opt.kernel_checks
        rule_friction = opt.rule_friction
    # Always persist the raw agent output, code optimization phase or not.
    _write_agent_transcripts(out_dir, translate, optimize)

    # Always emit a clean, marker-free production copy of the final solver.
    production_path = None
    if translate.success:
        production_path = write_production_variant(
            out_dir,
            target.solver_source_path(),
            tuple(rel for rel in editable if rel != target.solver_source_path()),
        )

    return CodegenResult(
        ir=ir,
        out_dir=out_dir,
        target=target,
        translate=translate,
        optimize=optimize,
        timing=timing,
        optimization_diff=optimization_diff,
        optimizations_md=optimizations_md,
        production_path=production_path,
        kernel_checks=kernel_checks,
        rule_friction=rule_friction,
    )

CodegenConfig

Bases: BaseModel

Top-level knobs controlling a code-generation run.

Attributes:

Name	Type	Description
version	Version	Embedded (no malloc / division-light) or large-scale.
backend	Backend	Target language. CUDA / OPENMP are only meaningful with version == LARGE_SCALE; C pairs with EMBEDDED.
termination	TerminationSpec	Loop termination criterion.
solve_method	SolveMethod	Linear-system solution strategy. AUTO selects an appropriate direct or iterative method from the version, backend, and problem size.
dtype	{'float64', 'float32'}	Floating-point type emitted in the C ABI ("float64" or "float32"). Large-scale GPU backends commonly prefer float32.
tol_verify	float or None	Tolerance used by the equivalence harness when comparing generated-code output against the Python reference oracle. When left unset it defaults by dtype (1e-6 for float64, 1e-4 for float32, since single precision cannot meet a 1e-6 gate); set it explicitly to override.
tol_kernel	float or None	Tolerance of the direct checks of each fom_apply_<name> and fom_fused_<block> against the package's reference, relative to 1 + |M x|. A rewritten kernel may only reorder arithmetic, so the float64 default (1e-10) is far tighter than tol_verify: a perturbed coefficient that the end-to-end check would absorb fails here. Defaults by dtype (1e-10 for float64, 1e-4 for float32).
num_verify_samples	int	Number of random parameter instances drawn for the golden verify data.
seed	int	RNG seed for reproducible verify-sample generation.
agent_timeout_s	int	Wall-clock budget, in seconds, for the translation phase, shared by all of its agent attempts. Default 600.
agent_optimize_max_turns	int	Turn budget for one coding-agent attempt at the code optimization phase. Default 200.
agent_optimize_timeout_s	int	Wall-clock budget, in seconds, for one round of the code optimization phase. Default 7200.
agent_optimize_rounds	int	Maximum number of code optimization rounds. Default 5.
agent_optimize_min_gain	float	A round that makes the solver less than this much faster (relative) ends the phase. Default 0.02.
cg_auto_min_dim	int	Dimension threshold for the AUTO → CG upgrade on the large-scale CUDA backend (see SolveMethod). Raise it to keep more systems on a direct factorization; lower it to prefer CG for smaller SPD systems. Ignored for non-CUDA backends and explicit solve methods.

Source code in fomtools/codegen/config.py

class CodegenConfig(BaseModel):
    """Top-level knobs controlling a code-generation run.

    Attributes
    ----------
    version : Version
        Embedded (no malloc / division-light) or large-scale.
    backend : Backend
        Target language. ``CUDA`` / ``OPENMP`` are only meaningful with
        ``version == LARGE_SCALE``; ``C`` pairs with ``EMBEDDED``.
    termination : TerminationSpec
        Loop termination criterion.
    solve_method : SolveMethod
        Linear-system solution strategy. ``AUTO`` selects an appropriate
        direct or iterative method from the version, backend, and problem size.
    dtype : {"float64", "float32"}
        Floating-point type emitted in the C ABI (``"float64"`` or
        ``"float32"``). Large-scale GPU backends commonly prefer ``float32``.
    tol_verify : float or None
        Tolerance used by the equivalence harness when comparing generated-code
        output against the Python reference oracle. When left unset it defaults
        by dtype (1e-6 for float64, 1e-4 for float32, since single precision
        cannot meet a 1e-6 gate); set it explicitly to override.
    tol_kernel : float or None
        Tolerance of the direct checks of each ``fom_apply_<name>`` and
        ``fom_fused_<block>`` against the package's reference, relative to
        ``1 + |M x|``. A rewritten kernel may only reorder arithmetic, so the
        float64 default (1e-10) is far tighter than ``tol_verify``: a perturbed
        coefficient that the end-to-end check would absorb fails here. Defaults
        by dtype (1e-10 for float64, 1e-4 for float32).
    num_verify_samples : int
        Number of random parameter instances drawn for the golden verify data.
    seed : int
        RNG seed for reproducible verify-sample generation.
    agent_timeout_s : int
        Wall-clock budget, in seconds, for the translation phase, shared by
        all of its agent attempts. Default 600.
    agent_optimize_max_turns : int
        Turn budget for one coding-agent attempt at the code optimization
        phase. Default 200.
    agent_optimize_timeout_s : int
        Wall-clock budget, in seconds, for one round of the code optimization
        phase. Default 7200.
    agent_optimize_rounds : int
        Maximum number of code optimization rounds. Default 5.
    agent_optimize_min_gain : float
        A round that makes the solver less than this much faster (relative)
        ends the phase. Default 0.02.
    cg_auto_min_dim : int
        Dimension threshold for the ``AUTO`` → CG upgrade on the large-scale
        CUDA backend (see `SolveMethod`). Raise it to keep more systems
        on a direct factorization; lower it to prefer CG for smaller SPD
        systems. Ignored for non-CUDA backends and explicit solve methods.
    """

    version: Version = Version.EMBEDDED
    backend: Backend = Backend.C
    termination: TerminationSpec = Field(default_factory=TerminationSpec)
    solve_method: SolveMethod = SolveMethod.AUTO
    dtype: Literal["float64", "float32"] = "float64"
    tol_verify: float | None = Field(
        default=None,
        description=(
            "Tolerance used by the equivalence harness when comparing generated "
            "code against the reference oracle. Defaults by dtype when left unset: "
            "1e-6 for float64, 1e-4 for float32 (single precision carries ~1e-7 "
            "relative error, which 1e-6 is too tight to absorb — especially for "
            "iterative CG). Set explicitly to override."
        ),
    )
    tol_kernel: float | None = Field(
        default=None,
        description=(
            "Tolerance of the direct kernel checks (fom_apply_*, fom_fused_*) "
            "against the package's reference, relative to 1 + |M x|. Defaults by "
            "dtype when left unset: 1e-10 for float64, 1e-4 for float32. A rewrite "
            "may only reorder arithmetic, so this is much tighter than tol_verify."
        ),
    )
    num_verify_samples: int = 8
    seed: int = 0
    agent_timeout_s: int = Field(
        ge=1,
        default=600,
        description=(
            "Wall-clock budget, in seconds, for the translation phase, shared by "
            "all of its agent attempts. The agent is told the time left after each "
            "tool call, a shell command cannot outlive it, and once it is spent the "
            "agent is stopped (after a short grace period to write its summary) "
            "and no further attempt starts. The phase is a bounded task, so the "
            "default is generous for it."
        ),
    )
    agent_optimize_max_turns: int = Field(
        ge=1,
        default=200,
        description=(
            "Turn budget for one coding-agent attempt at the code optimization phase. "
            "The phase is a profile/rewrite/verify/benchmark loop, so a method "
            "with more structure to exploit needs more turns; exhausting the "
            "budget leaves the verified solver in place but costs the agent's own "
            "summary of what it did. Sized so the wall-clock budget, not turns, "
            "is what ends a run at the default."
        ),
    )
    agent_optimize_timeout_s: int = Field(
        ge=1,
        default=7200,
        description=(
            "Wall-clock budget, in seconds, for one round of the code optimization "
            "phase (each round gets its own), enforced as for agent_timeout_s. Larger than the translation budget "
            "because the phase is open-ended: profile, rewrite, re-verify, "
            "benchmark, repeat, and each round of a GPU kernel rewrite (nvcc "
            "build, verify, benchmark) is slow. Two hours by default: a CUDA CG "
            "solve stopped at 30 minutes was still improving. A stopped agent "
            "still leaves whatever it had written on disk, which the pipeline "
            "re-verifies and benchmarks. The gates that run after the agent "
            "(verify, benchmark, profile) take seconds and are outside the budget."
        ),
    )
    agent_optimize_rounds: int = Field(
        ge=1,
        default=5,
        description=(
            "Maximum number of code optimization rounds. Each round is a fresh agent "
            "run on the fastest verified solver so far, told its current time and "
            "profile; the pipeline re-verifies and benchmarks it head to head "
            "against that solver, and keeps it only if it is faster. The phase "
            "ends early when a round fails its gates, is not faster, or gains less "
            "than agent_optimize_min_gain. One round is the single-pass behavior."
        ),
    )
    agent_optimize_min_gain: float = Field(
        ge=0.0,
        default=0.02,
        description=(
            "Relative speedup below which a kept round ends the code optimization "
            "phase: a round taking the solver from t to t' continues only if "
            "(t - t') / t >= agent_optimize_min_gain."
        ),
    )
    cg_auto_min_dim: int = Field(
        ge=1,
        default=256,
        description=(
            "Minimum affine-system dimension at which solve_method=AUTO upgrades "
            "a large-scale CUDA SPD solve to CG (matrix-free, data-parallel). "
            "Smaller SPD systems and all non-SPD systems keep a direct "
            "factorization. No effect for non-CUDA backends or an explicit "
            "solve_method."
        ),
    )

    model_config = ConfigDict(frozen=True)

    @model_validator(mode="before")
    @classmethod
    def _default_tol_verify(cls, data):
        # Fill the dtype-dependent tol_verify default when the user left it unset.
        if isinstance(data, dict) and data.get("tol_verify") is None:
            dtype = data.get("dtype", "float64")
            data["tol_verify"] = _TOL_VERIFY_DEFAULTS.get(dtype, 1e-6)
        if isinstance(data, dict) and data.get("tol_kernel") is None:
            dtype = data.get("dtype", "float64")
            data["tol_kernel"] = _TOL_KERNEL_DEFAULTS.get(dtype, 1e-10)
        return data

    @model_validator(mode="after")
    def _check_tol_verify(self) -> "CodegenConfig":
        if self.tol_verify is None or self.tol_verify <= 0:
            raise ValueError("tol_verify must be a positive float")
        if self.tol_kernel is None or self.tol_kernel <= 0:
            raise ValueError("tol_kernel must be a positive float")
        return self

    @model_validator(mode="after")
    def _check_version_constraints(self) -> "CodegenConfig":
        if self.version == Version.EMBEDDED and self.backend != Backend.C:
            raise ValueError(
                f"embedded version only supports the C backend; got backend={self.backend.value!r}"
            )
        if self.solve_method == SolveMethod.CG and self.version == Version.EMBEDDED:
            raise ValueError(
                "the CG (indirect) solve method needs runtime division and is only available "
                "for the large-scale version; the embedded target is division-free. Use "
                "solve_method=AUTO/LDLT/LU for the embedded version."
            )
        return self

    @property
    def c_float(self) -> str:
        """C type for the emitted ``fom_float`` typedef."""
        return "double" if self.dtype == "float64" else "float"

c_float property

c_float: str

C type for the emitted fom_float typedef.

TerminationSpec

Bases: BaseModel

Termination criterion for the generated solve loop.

Attributes:

Name	Type	Description
kind	TerminationKind	Fixed iteration count, infinity-norm change, or the verification objective, each with early exit except the first.
max_iters	int	Hard upper bound on iterations. Always required: embedded targets size their (static) iteration accounting from it and every loop needs a guard.
tol	float or None	Tolerance for the INF_NORM_CHANGE and OBJECTIVE criteria. Ignored for MAX_ITERS. Must be positive when the criterion uses it. For OBJECTIVE it bounds the verification objective as the verifier defines it, so with a 2-norm objective (ord=2) it bounds the SQUARED norm: tol=1e-6 stops once ||r||_2 <= 1e-3, not 1e-6.
check_termination	int	How often the criterion is evaluated (the name and semantics follow OSQP's setting): N checks every N iterations (residual computed only then), 0 disables early exit. Ignored for MAX_ITERS. For INF_NORM_CHANGE the default is 1 because the generated check is an infinity norm fused into the state copy (measured 1-2.5% of iteration time for direct-method FOMs), not extra residual matvecs; raise it for matrix-free methods on large problems where iterations are cheap. With 0 the loop runs max_iters (equivalent to MAX_ITERS but keeps the criterion recorded). For OBJECTIVE the default is 25: the objective costs its own matrix-vector products. The objective is also evaluated and tested at the last iteration (k == max_iter), whatever N is, so the reported residual is always the objective at the returned iterate and a solver run to a verified horizon reports FOM_SOLVED. With 0 it is evaluated only there.

Source code in fomtools/codegen/config.py

class TerminationSpec(BaseModel):
    """Termination criterion for the generated solve loop.

    Attributes
    ----------
    kind : TerminationKind
        Fixed iteration count, infinity-norm change, or the verification
        objective, each with early exit except the first.
    max_iters : int
        Hard upper bound on iterations. Always required: embedded targets size
        their (static) iteration accounting from it and every loop needs a guard.
    tol : float or None
        Tolerance for the ``INF_NORM_CHANGE`` and ``OBJECTIVE`` criteria. Ignored
        for ``MAX_ITERS``. Must be positive when the criterion uses it. For
        ``OBJECTIVE`` it bounds the verification objective as the verifier defines
        it, so with a 2-norm objective (``ord=2``) it bounds the SQUARED norm:
        ``tol=1e-6`` stops once ``||r||_2 <= 1e-3``, not ``1e-6``.
    check_termination : int
        How often the criterion is evaluated (the name and semantics follow
        OSQP's setting): ``N`` checks every ``N`` iterations (residual computed
        only then), ``0`` disables early exit. Ignored for ``MAX_ITERS``.

        For ``INF_NORM_CHANGE`` the default is 1 because the generated check is an
        infinity norm fused into the state copy (measured 1-2.5% of iteration
        time for direct-method FOMs), not extra residual matvecs; raise it for
        matrix-free methods on large problems where iterations are cheap. With
        ``0`` the loop runs ``max_iters`` (equivalent to ``MAX_ITERS`` but keeps
        the criterion recorded).

        For ``OBJECTIVE`` the default is 25: the objective costs its own
        matrix-vector products. The objective is also evaluated and tested at the
        last iteration (``k == max_iter``), whatever ``N`` is, so the reported
        residual is always the objective at the returned iterate and a solver run
        to a verified horizon reports ``FOM_SOLVED``. With ``0`` it is evaluated
        only there.
    """

    model_config = ConfigDict(frozen=True)

    kind: TerminationKind = TerminationKind.MAX_ITERS
    max_iters: int = Field(ge=1, default=1000, description="Hard upper bound on iterations.")
    tol: float | None = Field(
        default=None,
        description="Tolerance for the ``INF_NORM_CHANGE`` and ``OBJECTIVE`` criteria.",
    )
    check_termination: int = Field(
        ge=0,
        default=1,
        description=(
            "How often the termination criterion is evaluated (default 1 for "
            "``INF_NORM_CHANGE``, 25 for ``OBJECTIVE``)."
        ),
    )

    @model_validator(mode="before")
    @classmethod
    def _default_check_termination(cls, data):
        # The default interval depends on the criterion, so fill it here rather
        # than in the field.
        if isinstance(data, dict) and data.get("check_termination") is None:
            data = dict(data)
            data.pop("check_termination", None)
            kind = TerminationKind(data.get("kind", TerminationKind.MAX_ITERS))
            if kind.value in _CHECK_TERMINATION_DEFAULTS:
                data["check_termination"] = _CHECK_TERMINATION_DEFAULTS[kind.value]
        return data

    @model_validator(mode="after")
    def _check_tol(self) -> "TerminationSpec":
        if self.kind in _TOL_KINDS and (self.tol is None or self.tol <= 0):
            raise ValueError(f"tol must be a positive float for the {self.kind.value} criterion")
        return self

    @property
    def uses_tolerance(self) -> bool:
        """Whether the loop checks the infinity-norm change criterion at all.

        Single source of truth shared by the emitted header (``FOM_USE_TOL``)
        and the reference oracle, so the generated code and the verification
        ground truth can never disagree on when early exit is possible.
        """
        return self.kind == TerminationKind.INF_NORM_CHANGE and self.check_termination > 0

    @property
    def uses_objective(self) -> bool:
        """Whether the loop evaluates the verification objective (``FOM_USE_OBJECTIVE``).

        True even with ``check_termination == 0``: the objective is then still
        evaluated, and tested, once at the last iteration.
        """
        return self.kind == TerminationKind.OBJECTIVE

    def is_check_iteration(self, k: int, max_iter: int | None = None) -> bool:
        """Whether the 1-based iteration ``k`` evaluates the ``OBJECTIVE`` criterion.

        Every ``check_termination``-th iteration, and the last one (``max_iter``,
        defaulting to ``max_iters``). Shared by the oracle and the emitted
        verify data so both test at the same iterations.
        """
        last = self.max_iters if max_iter is None else max_iter
        check = self.check_termination
        return k == last or (check > 0 and k % check == 0)

uses_objective property

uses_objective: bool

Whether the loop evaluates the verification objective (FOM_USE_OBJECTIVE).

True even with check_termination == 0: the objective is then still evaluated, and tested, once at the last iteration.

uses_tolerance property

uses_tolerance: bool

Whether the loop checks the infinity-norm change criterion at all.

Single source of truth shared by the emitted header (FOM_USE_TOL) and the reference oracle, so the generated code and the verification ground truth can never disagree on when early exit is possible.

is_check_iteration

is_check_iteration(k: int, max_iter: int | None = None) -> bool

Whether the 1-based iteration k evaluates the OBJECTIVE criterion.

Every check_termination-th iteration, and the last one (max_iter, defaulting to max_iters). Shared by the oracle and the emitted verify data so both test at the same iterations.

Source code in fomtools/codegen/config.py

def is_check_iteration(self, k: int, max_iter: int | None = None) -> bool:
    """Whether the 1-based iteration ``k`` evaluates the ``OBJECTIVE`` criterion.

    Every ``check_termination``-th iteration, and the last one (``max_iter``,
    defaulting to ``max_iters``). Shared by the oracle and the emitted
    verify data so both test at the same iterations.
    """
    last = self.max_iters if max_iter is None else max_iter
    check = self.check_termination
    return k == last or (check > 0 and k % check == 0)

TerminationKind

Bases: str, Enum

How the generated solve loop decides to stop.

MAX_ITERS runs a fixed iteration count (exactly matches the verifier / simulator semantics, so it is the most robust to verify). INF_NORM_CHANGE additionally stops early once ||x_k - x_{k-1}||_inf <= tol.

OBJECTIVE stops early once the verification objective r^k -- the performance metric set with Verifier.maximize, e.g. a fixed-point or KKT residual, in the norm it was set with -- is at most tol. The stopping test is then exactly the certified quantity: if verification shows r^K <= tol over the whole parameter set, the generated solver stops within K iterations. The norm follows the verifier's convention, so ord=2 is the squared two-norm and tol bounds it.

Source code in fomtools/codegen/config.py

class TerminationKind(str, Enum):
    """How the generated solve loop decides to stop.

    ``MAX_ITERS`` runs a fixed iteration count (exactly matches the verifier /
    simulator semantics, so it is the most robust to verify). ``INF_NORM_CHANGE``
    additionally stops early once ``||x_k - x_{k-1}||_inf <= tol``.

    ``OBJECTIVE`` stops early once the *verification objective* ``r^k`` -- the
    performance metric set with `Verifier.maximize`,
    e.g. a fixed-point or KKT residual, in the norm it was set with -- is at most
    ``tol``. The stopping test is then exactly the certified quantity: if
    verification shows ``r^K <= tol`` over the whole parameter set, the generated
    solver stops within ``K`` iterations. The norm follows the verifier's
    convention, so ``ord=2`` is the *squared* two-norm and ``tol`` bounds it.
    """

    MAX_ITERS = "max_iters"
    INF_NORM_CHANGE = "inf_norm_change"
    OBJECTIVE = "objective"

extract_unrolled_ir

extract_unrolled_ir(verifier, iterates, config: CodegenConfig | None = None, name: str = 'fom', state=None, params=None, dedup: bool = True, check: bool = True, roll: bool = True, param_values: dict | None = None) -> FomIR

Extract the whole recorded run as one unrolled body.

Parameters:

Name	Type	Description	Default
verifier		The Verifier holding the recorded steps.	required
iterates		The iterate sequence, either an IterateTrajectory or a plain sequence of expressions. A leading entry that is the initial iterate is recognized and dropped, so both conventions work.	required
config	CodegenConfig | None	Code-generation configuration. Its termination is replaced by the recorded horizon (a single pass when unrolled). OBJECTIVE keeps its tolerance and check interval: a rolled result stops on the verification objective within the horizon, an unrolled one evaluates it once at the end. INF_NORM_CHANGE is rejected.	None
state		Optional explicit leaf lists, as for extract_ir.	None
params		Optional explicit leaf lists, as for extract_ir.	None
dedup	bool	Merge coefficient blocks holding identical numbers.	True
check	bool	Verify the unrolled result against the recorded tape before returning. Does not apply to a rolled result: there the check is what decides to roll in the first place, so it cannot be skipped without changing the answer. Combine with roll=False to return an IR that has not been verified.	True
roll	bool	Emit the compact rolled IR when repeating one iteration is verified to reproduce the recording. Pass False to unroll regardless.	True
param_values	dict | None	Parameter values for the checks, keyed by parameter leaf. Defaults to deterministic draws inside the declared box. Supplying them adds a draw to the roll decision rather than replacing the default ones: no single draw, the caller's included, is enough to license a roll.	None

Returns:

Type	Description
FomIR	Either a rolled IR (unrolled false, the recorded horizon in termination.max_iters), or an unrolled one whose body holds the whole run in dependency order with state_update producing the answer. unrolled says which.

Source code in fomtools/codegen/extract_unrolled.py

def extract_unrolled_ir(
    verifier,
    iterates,
    config: CodegenConfig | None = None,
    name: str = "fom",
    state=None,
    params=None,
    dedup: bool = True,
    check: bool = True,
    roll: bool = True,
    param_values: dict | None = None,
) -> FomIR:
    """Extract the whole recorded run as one unrolled body.

    Parameters
    ----------
    verifier:
        The ``Verifier`` holding the recorded steps.
    iterates:
        The iterate sequence, either an
        `IterateTrajectory` or a plain
        sequence of expressions. A leading entry that is the *initial* iterate is
        recognized and dropped, so both conventions work.
    config:
        Code-generation configuration. Its termination is replaced by the
        recorded horizon (a single pass when unrolled). ``OBJECTIVE`` keeps its
        tolerance and check interval: a rolled result stops on the verification
        objective within the horizon, an unrolled one evaluates it once at the
        end. ``INF_NORM_CHANGE`` is rejected.
    state, params:
        Optional explicit leaf lists, as for
        `extract_ir`.
    dedup:
        Merge coefficient blocks holding identical numbers.
    check:
        Verify the *unrolled* result against the recorded tape before returning.
        Does not apply to a rolled result: there the check is what decides to roll
        in the first place, so it cannot be skipped without changing the answer.
        Combine with ``roll=False`` to return an IR that has not been verified.
    roll:
        Emit the compact rolled IR when repeating one iteration is verified to
        reproduce the recording. Pass ``False`` to unroll regardless.
    param_values:
        Parameter values for the checks, keyed by parameter leaf. Defaults to
        deterministic draws inside the declared box. Supplying them *adds* a
        draw to the roll decision rather than replacing the default ones: no
        single draw, the caller's included, is enough to license a roll.

    Returns
    -------
    FomIR
        Either a rolled IR (``unrolled`` false, the recorded horizon in
        ``termination.max_iters``), or an unrolled one whose body holds the whole
        run in dependency order with ``state_update`` producing the answer.
        `unrolled` says which.
    """
    config = config or CodegenConfig()
    _require_supported_termination(config)

    groups = normalize_iterates(verifier, iterates)
    if not groups:
        raise ValueError(
            "iterates must contain at least one recorded iteration. Pass the iterate "
            "sequence the loop produced, or an IterateTrajectory."
        )

    # Try rolling before unrolling: for a stationary recording this is both the
    # better artifact and the cheaper path, since it never builds the big body.
    if roll:
        rolled = _try_roll(verifier, groups, config, name, state, params, param_values)
        if rolled is not None:
            log.info(
                "%s: repeating one iteration reproduces the %d recorded iteration(s); "
                "emitting a rolled body of %d step(s) instead of unrolling.",
                name,
                len(groups),
                len(rolled.body),
            )
            return rolled

    ir = extract_ir(
        verifier,
        groups[-1],
        state=state,
        params=params,
        config=config.model_copy(update={"termination": _horizon(config, 1)}),
        name=name,
        iterates=groups,
    )
    _check_no_blowup(ir, len(groups))
    log.info(
        "%s: the recording does not roll; emitting %d unrolled step(s) over "
        "%d recorded iteration(s).",
        name,
        len(ir.body),
        len(groups),
    )

    ir = replace(ir, unrolled=True)
    if dedup:
        ir = dedup_matrices(ir)
    if check:
        check_unrolled(ir, verifier, groups, param_values=param_values, state=state, params=params)
    return ir

extract_ir

extract_ir(verifier, next_state, state=None, params=None, config: CodegenConfig | None = None, name: str = 'fom', iterates=None) -> FomIR

Build a FomIR for one iteration of the FOM recorded in verifier.

Parameters:

Name	Type	Description	Default
verifier		A Verifier whose step API was used to build at least one iteration of the algorithm.	required
next_state		The iterate LinearExpression (or list thereof) produced after exactly one iteration. Its concatenation defines the loop-carried state update.	required
state		Optional explicit list of initial-state leaves. When None, inferred as the external leaves that appear in verifier.iterates.	None
params		Optional explicit list of parameter leaves. When None, inferred as the external leaves not in verifier.iterates.	None
config	CodegenConfig | None	Code-generation configuration (supplies the termination spec / dtype).	None
name	str	Human-friendly FOM name used for file naming.	'fom'
iterates		Optional recorded iterate sequence. Used only under TerminationKind.OBJECTIVE, to locate an objective written at the last recorded iteration (see fomtools.codegen.objective).	None

Returns:

Type	Description
FomIR	The single-iteration intermediate representation.

Raises:

Type	Description
ObjectiveExtractionError	Under TerminationKind.OBJECTIVE, if no objective is set on the verifier or it reads something the generated loop does not carry.

Source code in fomtools/codegen/extract.py

def extract_ir(
    verifier,
    next_state,
    state=None,
    params=None,
    config: CodegenConfig | None = None,
    name: str = "fom",
    iterates=None,
) -> FomIR:
    """Build a `FomIR` for one iteration of the FOM recorded in ``verifier``.

    Parameters
    ----------
    verifier:
        A ``Verifier`` whose ``step`` API was used to build at least one
        iteration of the algorithm.
    next_state:
        The iterate ``LinearExpression`` (or list thereof) produced after exactly
        one iteration. Its concatenation defines the loop-carried state update.
    state:
        Optional explicit list of initial-state leaves. When ``None``, inferred
        as the external leaves that appear in ``verifier.iterates``.
    params:
        Optional explicit list of parameter leaves. When ``None``, inferred as
        the external leaves not in ``verifier.iterates``.
    config:
        Code-generation configuration (supplies the termination spec / dtype).
    name:
        Human-friendly FOM name used for file naming.
    iterates:
        Optional recorded iterate sequence. Used only under
        ``TerminationKind.OBJECTIVE``, to locate an objective written at the last
        recorded iteration (see `fomtools.codegen.objective`).

    Returns
    -------
    FomIR
        The single-iteration intermediate representation.

    Raises
    ------
    ObjectiveExtractionError
        Under ``TerminationKind.OBJECTIVE``, if no objective is set on the
        verifier or it reads something the generated loop does not carry.
    """
    config = config or CodegenConfig()
    next_state_exprs, body_steps, body_fresh, state_leaves, param_leaves = resolve_structure(
        verifier, next_state, state, params
    )

    # --- symbolic environment -------------------------------------------------
    state_layout = []  # (leaf, slot_name, dim, offset)
    offset = 0
    for j, leaf in enumerate(state_leaves):
        state_layout.append((leaf, f"s{j}", leaf.ndim, offset))
        offset += leaf.ndim
    state_dim = offset

    param_map = {leaf: (f"p{i}", leaf.ndim) for i, leaf in enumerate(param_leaves)}

    matrices: dict[str, MatrixData] = {}
    factorizations: dict[str, Factorization] = {}
    body: list[IterationStep] = []
    tmp_map: dict[LinearExpression, tuple[str, int]] = {}
    # Affine steps sharing an operator share one factorization. An unrolled
    # recording repeats the same solve every iteration, so this saves the
    # extraction-time factorization work as well as the emitted data.
    fac_cache: dict[tuple, str] = {}

    for i, step in enumerate(body_steps):
        source, activation, out_dim = _step_source_activation(step)
        prefix = f"t{i}"
        lin = _build_linear_op(source, prefix, state_layout, param_map, tmp_map, matrices, out_dim)
        solve_name = None
        # Affine operator: factorize M once (never inverted); the solver applies
        # division-free triangular solves to the assembled right-hand side.
        if isinstance(step, Affine) and not _is_identity(step.output_mat):
            method = _resolve_solve_method(step.output_mat, config)
            key = (_matrix_key(step.output_mat), method)
            solve_name = fac_cache.get(key)
            if solve_name is None:
                solve_name = fac_cache[key] = prefix
                factorizations[solve_name] = factorize(
                    step.output_mat, method=method, cg_tol=_cg_tol_for_dtype(config.dtype)
                )
        # ``M w = rhs`` with ``M = I`` is a copy, not a solve, and keeps its
        # identity activation. Trajectory barriers
        # (``IterateTrajectory._make_linked_leaf``) record exactly this, so
        # factorizing here would cost a stored LDL^T plus two triangular solves
        # per iteration for nothing.
        # Saturated-linear clip bounds become per-coordinate IR arrays.
        if activation.kind == ActivationKind.SATURATED_LINEAR:
            lower, upper = _satlin_arrays(step)
            lo_name, hi_name = f"{prefix}_lo", f"{prefix}_hi"
            matrices[lo_name] = MatrixData(name=lo_name, array=lower)
            matrices[hi_name] = MatrixData(name=hi_name, array=upper)
            activation = ActivationOp(
                kind=ActivationKind.SATURATED_LINEAR, lower=lo_name, upper=hi_name
            )
        body.append(
            IterationStep(
                name=prefix, out_dim=out_dim, linear_op=lin, activation=activation, solve=solve_name
            )
        )
        tmp_map[body_fresh[i]] = (prefix, out_dim)

    # --- state update ---------------------------------------------------------
    state_update = _build_state_update(
        next_state_exprs, state_layout, param_map, tmp_map, matrices, state_dim
    )
    if state_update.out_dim != state_dim:
        raise NonStationaryFOMError(
            f"next_state dimension ({state_update.out_dim}) does not match the "
            f"inferred state dimension ({state_dim})."
        )

    # --- verification objective (early exit on the certified metric) ----------
    objective = None
    if config.termination.uses_objective:
        from fomtools.codegen.objective import extract_objective

        objective = extract_objective(
            verifier,
            state_layout,
            next_state_exprs,
            param_map,
            tmp_map,
            matrices,
            iterates=normalize_iterates(verifier, iterates) if iterates is not None else None,
        )

    # --- initial state + param specs -----------------------------------------
    init = _initial_state(state_layout)
    state_spec = StateSpec(
        dim=state_dim,
        slots=tuple((slot_name, dim) for (_, slot_name, dim, _) in state_layout),
        init=init,
    )
    param_specs = tuple(
        ParamSpec(
            name=param_map[leaf][0],
            dim=leaf.ndim,
            lb=np.asarray(leaf.lb, dtype=float).reshape(-1),
            ub=np.asarray(leaf.ub, dtype=float).reshape(-1),
        )
        for leaf in param_leaves
    )

    return FomIR(
        params=param_specs,
        state=state_spec,
        body=tuple(body),
        state_update=state_update,
        termination=config.termination,
        matrices=matrices,
        factorizations=factorizations,
        dtype=config.dtype,
        name=name,
        objective=objective,
    )

extract_scheduled_ir

extract_scheduled_ir(build: Callable[..., Recording], schedules: dict, config: CodegenConfig | None = None, name: str = 'fom') -> FomIR

Extract an IR whose coefficients vary with the iteration index.

Parameters:

Name	Type	Description	Default
build	Callable[..., Recording]	A function that records one iteration and returns a Recording. It is called with one keyword argument per entry of schedules, and must be free of side effects: it is called several times, at different values of those scalars, and must build the same algorithm each time.	required
schedules	dict	Maps a schedule name to a callable of the 0-based iteration index, a sequence of values, or a Schedule. A textbook formula written in 1-based k must be shifted by the caller; the coefficient belongs to the method, not to the generator.	required
config	CodegenConfig | None	Code-generation configuration; its termination.max_iters fixes the length of every emitted table.	None
name	str	Human-friendly FOM name used for file naming.	'fom'

Returns:

Type	Description
FomIR	A single-iteration IR carrying the schedule tables and scaled terms.

Raises:

Type	Description
NonStationaryFOMError	If the recording changes shape between probes.
ScheduleFitError	If a coefficient's dependence on the schedule is not affine, or the varying quantity reaches somewhere a table cannot (a factorized system matrix, an activation scalar, a constant offset).

Source code in fomtools/codegen/extract_scheduled.py

def extract_scheduled_ir(
    build: Callable[..., Recording],
    schedules: dict,
    config: CodegenConfig | None = None,
    name: str = "fom",
) -> FomIR:
    """Extract an IR whose coefficients vary with the iteration index.

    Parameters
    ----------
    build:
        A function that records **one** iteration and returns a
        `Recording`. It is called with one keyword argument per entry of
        ``schedules``, and must be free of side effects: it is called several
        times, at different values of those scalars, and must build the same
        algorithm each time.
    schedules:
        Maps a schedule name to a callable of the **0-based** iteration index, a
        sequence of values, or a `Schedule`. A
        textbook formula written in 1-based ``k`` must be shifted by the caller;
        the coefficient belongs to the method, not to the generator.
    config:
        Code-generation configuration; its ``termination.max_iters`` fixes the
        length of every emitted table.
    name:
        Human-friendly FOM name used for file naming.

    Returns
    -------
    FomIR
        A single-iteration IR carrying the schedule tables and scaled terms.

    Raises
    ------
    NonStationaryFOMError
        If the recording changes shape between probes.
    ScheduleFitError
        If a coefficient's dependence on the schedule is not affine, or the
        varying quantity reaches somewhere a table cannot (a factorized system
        matrix, an activation scalar, a constant offset).
    """
    config = config or CodegenConfig()
    if not schedules:
        raise ValueError(
            "extract_scheduled_ir() needs at least one schedule; use extract_ir() "
            "for a stationary method."
        )
    names = sorted(schedules)

    tables = {n: as_schedule(n, schedules[n], config.termination.max_iters) for n in names}

    # One probe at the base point, one per schedule with that scalar moved, and a
    # final all-moved point held out to validate the recovered fit.
    base_point = {n: _PROBE_BASE for n in names}
    points = [base_point]
    for n in names:
        moved = dict(base_point)
        moved[n] = _PROBE_STEP
        points.append(moved)
    points.append({n: _PROBE_CHECK for n in names})

    irs = [_extract_at(build, point, config, name) for point in points]

    _check_same_structure(irs, points)
    _check_scalars_fixed(irs)
    _check_factorizations_fixed(irs)

    bases, slopes = _recover_slopes(irs, names)

    ref = irs[0]
    bound_names = set()
    for step in ref.body:
        for attr in ("lower", "upper"):
            bound = getattr(step.activation, attr)
            if bound is not None:
                bound_names.add(bound)
    for mname in slopes:
        if mname in bound_names:
            raise ScheduleFitError(
                f"the saturated-linear bound {mname!r} varies with the iteration index. "
                "Scheduled clip bounds are not supported yet; only linear-operator "
                "coefficients can be scheduled."
            )

    matrices: dict[str, MatrixData] = {
        mname: MatrixData(
            name=mname,
            array=bases[mname],
            sparse_pattern=ref.matrices[mname].sparse_pattern if mname not in slopes else None,
        )
        for mname in ref.matrices
    }
    for mname, per_schedule in slopes.items():
        for sched, slope in per_schedule.items():
            dname = f"{mname}__d_{sched}"
            matrices[dname] = MatrixData(name=dname, array=slope)

    body = tuple(
        replace(step, linear_op=_rewrite_linop(step.linear_op, slopes)) for step in ref.body
    )
    state_update = _rewrite_linop(ref.state_update, slopes)
    # The objective is evaluated on the iterate an iteration produced, so a term
    # that varies with the schedule reads the same table entry as that iteration.
    objective = None
    if ref.objective is not None:
        objective = replace(
            ref.objective,
            blocks=tuple(
                replace(block, linear_op=_rewrite_linop(block.linear_op, slopes))
                for block in ref.objective.blocks
            ),
        )

    used = {
        t.scale
        for lin in [*(s.linear_op for s in body), state_update]
        + (objective.linops() if objective is not None else [])
        for t in lin.terms
    }
    unused = [n for n in names if n not in used]
    if unused:
        raise ScheduleFitError(
            f"schedule(s) {unused!r} do not affect any coefficient. The build function "
            "must actually use every scalar it is given; an unused schedule usually "
            "means a typo in its keyword name."
        )

    return FomIR(
        params=ref.params,
        state=ref.state,
        body=body,
        state_update=state_update,
        termination=ref.termination,
        matrices=matrices,
        factorizations=ref.factorizations,
        schedules={n: tables[n].values for n in names},
        dtype=ref.dtype,
        name=name,
        objective=objective,
    )

FomIR dataclass

Complete IR for one repeatable iteration, or for a whole unrolled run.

Attributes:

Name	Type	Description
params	tuple[ParamSpec, ...]	Solve-time parameter vectors.
state	StateSpec	Loop-carried state layout + initial value. For an unrolled IR this is the solver's input: the initial iterates, plus any earlier iterate the method keeps referring to (a Halpern anchor, say).
body	tuple[IterationStep, ...]	Ordered iteration-body steps producing temporaries. For an unrolled IR this is the whole run as one straight-line program, in dependency order.
state_update	LinearOp	Linear op producing the next state from current state + temporaries. For an unrolled IR it produces the final answer.
termination	TerminationSpec	Loop termination criterion.
matrices	dict[str, MatrixData]	All named coefficient/const blocks referenced by the ops.
schedules	dict[str, ndarray]	Iteration-varying scalars referenced by Term.scale, each a table of at least termination.max_iters values indexed by the 0-based iteration counter. Empty for a fully stationary method. Mutually exclusive with unrolled: both express iteration-varying coefficients, and they do it in incompatible ways.
dtype	str	Numeric dtype string ("float64" / "float32").
name	str	Human-friendly FOM name (used for fixtures / file naming).
unrolled	bool	Whether body is a whole run rather than one repeatable iteration. An unrolled body is executed once, so termination.max_iters is 1 and no residual-based early exit applies; each iteration carries its own coefficient blocks, which is what lets a recording use arbitrary per-iteration steps.
objective	ObjectiveSpec | None	The verification objective the loop stops on under OBJECTIVE, or None. For an unrolled IR it is evaluated once, after the single pass.

Source code in fomtools/codegen/ir.py

@dataclass(frozen=True)
class FomIR:
    """Complete IR for one repeatable iteration, or for a whole unrolled run.

    Attributes
    ----------
    params:
        Solve-time parameter vectors.
    state:
        Loop-carried state layout + initial value. For an unrolled IR this is
        the solver's *input*: the initial iterates, plus any earlier iterate the
        method keeps referring to (a Halpern anchor, say).
    body:
        Ordered iteration-body steps producing temporaries. For an unrolled IR
        this is the whole run as one straight-line program, in dependency order.
    state_update:
        Linear op producing the next state from current state + temporaries. For
        an unrolled IR it produces the final answer.
    termination:
        Loop termination criterion.
    matrices:
        All named coefficient/const blocks referenced by the ops.
    schedules:
        Iteration-varying scalars referenced by `Term.scale`, each a table
        of at least ``termination.max_iters`` values indexed by the 0-based
        iteration counter. Empty for a fully stationary method. Mutually
        exclusive with `unrolled`: both express iteration-varying
        coefficients, and they do it in incompatible ways.
    dtype:
        Numeric dtype string (``"float64"`` / ``"float32"``).
    name:
        Human-friendly FOM name (used for fixtures / file naming).
    unrolled:
        Whether `body` is a whole run rather than one repeatable
        iteration. An unrolled body is executed **once**, so
        ``termination.max_iters`` is 1 and no residual-based early exit applies;
        each iteration carries its own coefficient blocks, which is what lets a
        recording use arbitrary per-iteration steps.
    objective:
        The verification objective the loop stops on under
        `OBJECTIVE`, or ``None``.
        For an unrolled IR it is evaluated once, after the single pass.
    """

    params: tuple[ParamSpec, ...]
    state: StateSpec
    body: tuple[IterationStep, ...]
    state_update: LinearOp
    termination: TerminationSpec
    matrices: dict[str, MatrixData]
    factorizations: dict[str, Factorization] = field(default_factory=dict)
    schedules: dict[str, np.ndarray] = field(default_factory=dict)
    dtype: str = "float64"
    name: str = "fom"
    unrolled: bool = False
    objective: ObjectiveSpec | None = None

    def __post_init__(self):
        # Enforced here rather than left to "whichever extraction function you
        # happened to call": the two representations answer the same question
        # differently, and an IR holding both would make the emitted spec
        # self-contradictory -- "executed once, ignore max_iter" next to a table
        # indexed by a loop counter that no longer exists.
        if self.unrolled and self.schedules:
            raise ValueError(
                "an IR cannot be both unrolled and scheduled: an unrolled body already "
                "carries its own coefficients per iteration, so the schedule(s) "
                f"{sorted(self.schedules)} would be tabulated over a loop that is not "
                "emitted. Extract the run with extract_unrolled_ir (arbitrary variation, "
                "size grows with the horizon) or declare the variation with "
                "extract_scheduled_ir (compact, affine dependence only)."
            )
        if self.termination.uses_objective and self.objective is None:
            raise ValueError(
                "termination kind 'objective' needs the verification objective in the IR; "
                "set it on the verifier with maximize() before extracting."
            )

    @property
    def state_dim(self) -> int:
        return self.state.dim

    @property
    def is_scheduled(self) -> bool:
        """Whether any coefficient varies with the iteration index."""
        return bool(self.schedules)

    @property
    def max_rhs_dim(self) -> int:
        """Largest intermediate vector size; used to statically size scratch."""
        dims = [self.state.dim, self.state_update.out_dim]
        dims += [step.out_dim for step in self.body]
        dims += [step.linear_op.out_dim for step in self.body]
        return max(dims)

    def tmp_dims(self) -> dict[str, int]:
        """Map each temporary name to its dimension."""
        return {step.name: step.out_dim for step in self.body}

    def objective_inputs(self) -> list[str]:
        """The parameter and temporary sources the objective reads, in the order
        ``fom_objective`` takes them after ``(it, state, new_state)``: parameters in
        declaration order, then temporaries in body order."""
        if self.objective is None:
            return []
        used = {t.source for lin in self.objective.linops() for t in lin.terms}
        params = [f"param:{p.name}" for p in self.params if f"param:{p.name}" in used]
        tmps = [f"tmp:{s.name}" for s in self.body if f"tmp:{s.name}" in used]
        return params + tmps

    def linops(self) -> list[LinearOp]:
        """Every linear op: the body steps', the state update, and the objective's."""
        ops = [step.linear_op for step in self.body] + [self.state_update]
        if self.objective is not None:
            ops += self.objective.linops()
        return ops

is_scheduled property

is_scheduled: bool

Whether any coefficient varies with the iteration index.

max_rhs_dim property

max_rhs_dim: int

Largest intermediate vector size; used to statically size scratch.

linops

linops() -> list[LinearOp]

Every linear op: the body steps', the state update, and the objective's.

Source code in fomtools/codegen/ir.py

def linops(self) -> list[LinearOp]:
    """Every linear op: the body steps', the state update, and the objective's."""
    ops = [step.linear_op for step in self.body] + [self.state_update]
    if self.objective is not None:
        ops += self.objective.linops()
    return ops

objective_inputs

objective_inputs() -> list[str]

The parameter and temporary sources the objective reads, in the order fom_objective takes them after (it, state, new_state): parameters in declaration order, then temporaries in body order.

Source code in fomtools/codegen/ir.py

def objective_inputs(self) -> list[str]:
    """The parameter and temporary sources the objective reads, in the order
    ``fom_objective`` takes them after ``(it, state, new_state)``: parameters in
    declaration order, then temporaries in body order."""
    if self.objective is None:
        return []
    used = {t.source for lin in self.objective.linops() for t in lin.terms}
    params = [f"param:{p.name}" for p in self.params if f"param:{p.name}" in used]
    tmps = [f"tmp:{s.name}" for s in self.body if f"tmp:{s.name}" in used]
    return params + tmps

tmp_dims

tmp_dims() -> dict[str, int]

Map each temporary name to its dimension.

Source code in fomtools/codegen/ir.py

def tmp_dims(self) -> dict[str, int]:
    """Map each temporary name to its dimension."""
    return {step.name: step.out_dim for step in self.body}

extract_objective

extract_objective(verifier, state_layout, next_state_exprs, param_map, tmp_map, matrices: dict, iterates=None) -> ObjectiveSpec

Map the objective set on verifier onto the loop extracted by extract_ir.

Parameters:

Name	Type	Description	Default
verifier		The recording, with its objective set.	required
state_layout		extract_ir's view of the iteration: (leaf, slot, dim, offset) per state leaf, the iterate after one iteration, and the parameter / temporary leaf maps.	required
next_state_exprs		extract_ir's view of the iteration: (leaf, slot, dim, offset) per state leaf, the iterate after one iteration, and the parameter / temporary leaf maps.	required
param_map		extract_ir's view of the iteration: (leaf, slot, dim, offset) per state leaf, the iterate after one iteration, and the parameter / temporary leaf maps.	required
tmp_map		extract_ir's view of the iteration: (leaf, slot, dim, offset) per state leaf, the iterate after one iteration, and the parameter / temporary leaf maps.	required
matrices	dict	The IR's matrix table; the objective's coefficient blocks are added to it.	required
iterates		The recorded iterate sequence (one entry per iteration, normalized), if any. Enables the last-recorded-iteration frame.	None

Raises:

Type	Description
ObjectiveExtractionError	If no objective is set, or it reads something the loop does not carry.

Source code in fomtools/codegen/objective.py

def extract_objective(
    verifier,
    state_layout,
    next_state_exprs,
    param_map,
    tmp_map,
    matrices: dict,
    iterates=None,
) -> ObjectiveSpec:
    """Map the objective set on ``verifier`` onto the loop extracted by
    `extract_ir`.

    Parameters
    ----------
    verifier:
        The recording, with its objective set.
    state_layout, next_state_exprs, param_map, tmp_map:
        `extract_ir`'s view of the iteration:
        ``(leaf, slot, dim, offset)`` per state leaf, the iterate after one
        iteration, and the parameter / temporary leaf maps.
    matrices:
        The IR's matrix table; the objective's coefficient blocks are added to it.
    iterates:
        The recorded iterate sequence (one entry per iteration, normalized), if
        any. Enables the last-recorded-iteration frame.

    Raises
    ------
    ObjectiveExtractionError
        If no objective is set, or it reads something the loop does not carry.
    """
    if verifier.objective is None:
        raise ObjectiveExtractionError(
            "termination kind 'objective' stops on the verification objective, but none "
            "is set on this recording. Call verifier.maximize(ft.norm(expr, np.inf)) (e.g. "
            "for the fixed-point residual x_K - x_{K-1}) before generating code, "
            "or use TerminationKind.INF_NORM_CHANGE / MAX_ITERS."
        )
    exprs, ord = verifier.objective
    norm = norm_label(ord)
    state_dim = sum(dim for (_, _, dim, _) in state_layout)

    groups = [list(g) for g in iterates] if iterates else []
    loop_steps = _dependency_cone(
        verifier.steps,
        {leaf for g in [next_state_exprs, *groups] for e in g for leaf in e.decomposition_dict},
    )
    relu_steps = _relu_steps_outside(verifier, loop_steps)
    # Parameters are the leaves no step produces and no iterate declares, as in
    # extraction's own classification.
    produced = {leaf for step in verifier.steps for leaf in step.output_leaves}
    iterates_declared = set(verifier.iterates)
    outside_params = {
        leaf
        for expr in exprs
        for leaf in expr.decomposition_dict
        if leaf not in produced and leaf not in iterates_declared and leaf not in param_map
    }

    state_rows = {leaf: np.arange(off, off + dim) for (leaf, _, dim, off) in state_layout}
    frames = [
        _Frame(
            label="the extracted iteration (w^0 = the state leaves, w^1 = next_state)",
            state_dim=state_dim,
            new_rows=_copy_rows(next_state_exprs),
            prev_rows=state_rows,
            tmps=dict(tmp_map),
        )
    ]
    if len(groups) >= 2:
        K = len(groups)
        frames.append(
            _Frame(
                label=f"the last recorded iteration (w^{K - 1} and w^{K} from the tape)",
                state_dim=state_dim,
                new_rows=_copy_rows(groups[-1]),
                prev_rows=_copy_rows(groups[-2]),
            )
        )

    failures = []
    for frame in frames:
        try:
            mapper = _Mapper(frame, param_map, relu_steps, outside_params)
            blocks = _lower(exprs, mapper, matrices)
        except ObjectiveExtractionError as err:
            failures.append(str(err))
            continue
        return ObjectiveSpec(blocks=tuple(blocks), norm=norm)
    hint = (
        ""
        if len(frames) > 1
        else " To set the objective at the last recorded iteration instead, also pass "
        "iterates= (the recorded iterate sequence) so that iteration can be located."
    )
    raise ObjectiveExtractionError(
        "the verification objective cannot be evaluated by the generated loop.\n- "
        + "\n- ".join(failures)
        + hint
    )

ObjectiveExtractionError

Bases: ValueError

The verification objective cannot be evaluated by the generated loop.

Source code in fomtools/codegen/objective.py

class ObjectiveExtractionError(ValueError):
    """The verification objective cannot be evaluated by the generated loop."""

ReferenceOracle

Rolled NumPy executor of a FomIR.

Source code in fomtools/codegen/oracle.py

class ReferenceOracle:
    """Rolled NumPy executor of a `FomIR`."""

    def __init__(self, ir: FomIR):
        self.ir = ir

    def _eval_linop(self, lin: LinearOp, sources: dict[str, np.ndarray], it: int = 0) -> np.ndarray:
        if lin.const is not None:
            out = self.ir.matrices[lin.const].array.astype(float).reshape(-1).copy()
        else:
            out = np.zeros(lin.out_dim)
        for term in lin.terms:
            mat = self.ir.matrices[term.matrix].array
            contrib = np.asarray(mat @ sources[term.source]).reshape(-1)
            if term.scale is not None:
                # Scheduled term: the emitted C reads the same table at the same
                # 0-based loop counter.
                contrib = self.ir.schedules[term.scale][it] * contrib
            out = out + contrib
        return out

    def objective_value(self, sources: dict[str, np.ndarray], it: int) -> float:
        """The verification objective ``r^k`` at the end of 0-based iteration ``it``.

        ``sources`` must hold ``state`` (``w^{k-1}``), ``new_state`` (``w^k``), the
        parameters and that iteration's temporaries.
        """
        obj = self.ir.objective
        rows = [
            _apply_activation(
                block.activation, self._eval_linop(block.linear_op, sources, it), self.ir
            )
            for block in obj.blocks
        ]
        return objective_norm(np.concatenate(rows) if rows else np.zeros(0), obj.norm)

    def solve(
        self,
        param_values: dict[str, np.ndarray] | None = None,
        x0: np.ndarray | None = None,
    ) -> OracleResult:
        """Run the rolled loop.

        Parameters
        ----------
        param_values:
            Mapping from parameter name (``"p0"`` …) to its value vector.
        x0:
            Optional initial state; defaults to the IR's stored init.
        """
        ir = self.ir
        param_values = param_values or {}
        state = (
            np.array(ir.state.init, dtype=float)
            if x0 is None
            else np.asarray(x0, dtype=float).copy()
        )

        sources: dict[str, np.ndarray] = {"state": state}
        for p in ir.params:
            if p.name not in param_values:
                raise ValueError(f"Missing value for parameter {p.name!r}")
            sources[f"param:{p.name}"] = np.asarray(param_values[p.name], dtype=float).reshape(-1)

        trajectory = [state.copy()]
        residuals: list[float] = []
        objectives: list[tuple[int, float]] = []
        solved = False
        term = ir.termination
        # Early exit only on "check" iterations: every `check_termination`
        # iterations (0 disables checking), mirroring the generated C exactly.
        # The objective is also tested at the last iteration.
        check = term.check_termination
        use_tol = term.uses_tolerance
        use_objective = term.uses_objective
        iters = 0
        for it in range(term.max_iters):
            for step in ir.body:
                lin_val = self._eval_linop(step.linear_op, sources, it)
                if step.solve is not None:
                    # Affine step: solve M x = rhs via the factorization (the
                    # inverse is never formed), matching the emitted C exactly.
                    val = ir.factorizations[step.solve].solve(lin_val)
                else:
                    val = _apply_activation(step.activation, lin_val, ir)
                sources[f"tmp:{step.name}"] = val
            new_state = self._eval_linop(ir.state_update, sources, it)
            residual = float(np.max(np.abs(new_state - state))) if new_state.size else 0.0
            objective = None
            if use_objective and term.is_check_iteration(it + 1):
                # Before the state advances: the objective reads w^{k-1} and w^k.
                sources["new_state"] = new_state
                objective = self.objective_value(sources, it)
                objectives.append((it + 1, objective))
            state = new_state
            sources["state"] = state
            trajectory.append(state.copy())
            residuals.append(residual)
            iters += 1
            if use_tol and (it + 1) % check == 0 and residual <= term.tol:
                solved = True
                break
            if objective is not None and objective <= term.tol:
                solved = True
                break

        if use_objective:
            reported = objectives[-1][1] if objectives else 0.0
        else:
            reported = residuals[-1] if residuals else 0.0
        return OracleResult(
            solution=state,
            iters=iters,
            residual=reported,
            trajectory=trajectory,
            residuals=residuals,
            objectives=objectives,
            solved=solved,
        )

objective_value

objective_value(sources: dict[str, ndarray], it: int) -> float

The verification objective r^k at the end of 0-based iteration it.

sources must hold state (w^{k-1}), new_state (w^k), the parameters and that iteration's temporaries.

Source code in fomtools/codegen/oracle.py

def objective_value(self, sources: dict[str, np.ndarray], it: int) -> float:
    """The verification objective ``r^k`` at the end of 0-based iteration ``it``.

    ``sources`` must hold ``state`` (``w^{k-1}``), ``new_state`` (``w^k``), the
    parameters and that iteration's temporaries.
    """
    obj = self.ir.objective
    rows = [
        _apply_activation(
            block.activation, self._eval_linop(block.linear_op, sources, it), self.ir
        )
        for block in obj.blocks
    ]
    return objective_norm(np.concatenate(rows) if rows else np.zeros(0), obj.norm)

solve

solve(param_values: dict[str, ndarray] | None = None, x0: ndarray | None = None) -> OracleResult

Run the rolled loop.

Parameters:

Name	Type	Description	Default
param_values	dict[str, ndarray] | None	Mapping from parameter name ("p0" …) to its value vector.	None
x0	ndarray | None	Optional initial state; defaults to the IR's stored init.	None

Source code in fomtools/codegen/oracle.py

def solve(
    self,
    param_values: dict[str, np.ndarray] | None = None,
    x0: np.ndarray | None = None,
) -> OracleResult:
    """Run the rolled loop.

    Parameters
    ----------
    param_values:
        Mapping from parameter name (``"p0"`` …) to its value vector.
    x0:
        Optional initial state; defaults to the IR's stored init.
    """
    ir = self.ir
    param_values = param_values or {}
    state = (
        np.array(ir.state.init, dtype=float)
        if x0 is None
        else np.asarray(x0, dtype=float).copy()
    )

    sources: dict[str, np.ndarray] = {"state": state}
    for p in ir.params:
        if p.name not in param_values:
            raise ValueError(f"Missing value for parameter {p.name!r}")
        sources[f"param:{p.name}"] = np.asarray(param_values[p.name], dtype=float).reshape(-1)

    trajectory = [state.copy()]
    residuals: list[float] = []
    objectives: list[tuple[int, float]] = []
    solved = False
    term = ir.termination
    # Early exit only on "check" iterations: every `check_termination`
    # iterations (0 disables checking), mirroring the generated C exactly.
    # The objective is also tested at the last iteration.
    check = term.check_termination
    use_tol = term.uses_tolerance
    use_objective = term.uses_objective
    iters = 0
    for it in range(term.max_iters):
        for step in ir.body:
            lin_val = self._eval_linop(step.linear_op, sources, it)
            if step.solve is not None:
                # Affine step: solve M x = rhs via the factorization (the
                # inverse is never formed), matching the emitted C exactly.
                val = ir.factorizations[step.solve].solve(lin_val)
            else:
                val = _apply_activation(step.activation, lin_val, ir)
            sources[f"tmp:{step.name}"] = val
        new_state = self._eval_linop(ir.state_update, sources, it)
        residual = float(np.max(np.abs(new_state - state))) if new_state.size else 0.0
        objective = None
        if use_objective and term.is_check_iteration(it + 1):
            # Before the state advances: the objective reads w^{k-1} and w^k.
            sources["new_state"] = new_state
            objective = self.objective_value(sources, it)
            objectives.append((it + 1, objective))
        state = new_state
        sources["state"] = state
        trajectory.append(state.copy())
        residuals.append(residual)
        iters += 1
        if use_tol and (it + 1) % check == 0 and residual <= term.tol:
            solved = True
            break
        if objective is not None and objective <= term.tol:
            solved = True
            break

    if use_objective:
        reported = objectives[-1][1] if objectives else 0.0
    else:
        reported = residuals[-1] if residuals else 0.0
    return OracleResult(
        solution=state,
        iters=iters,
        residual=reported,
        trajectory=trajectory,
        residuals=residuals,
        objectives=objectives,
        solved=solved,
    )

AgentDriver

Bases: ABC

Abstract coding-agent that generates a source file until verify passes.

Subclasses implement run to drive a particular backend (the real Claude/Codex/Pi CLIs, or StubAgentDriver for tests). The driver owns only task.solver_rel_path; all other files are restored via task.restore_hook before each verify.

Methods:

Name	Description
run	Produce/refine the solver source until the verify command passes.

Source code in fomtools/codegen/agent/driver.py

class AgentDriver(ABC):
    """Abstract coding-agent that generates a source file until verify passes.

    Subclasses implement `run` to drive a particular backend (the real
    Claude/Codex/Pi CLIs, or `StubAgentDriver`
    for tests). The driver owns only ``task.solver_rel_path``; all other files
    are restored via ``task.restore_hook`` before each verify.

    Methods
    -------
    run
        Produce/refine the solver source until the verify command passes.
    """

    @abstractmethod
    def run(self, task: AgentTask) -> AgentResult:
        """Produce/refine ``task.solver_rel_path`` until ``task.verify_cmd`` passes.

        Parameters
        ----------
        task : AgentTask
            The unit of work describing the workspace, the solver path to
            write, the restore hook, and the verify command to run.

        Returns
        -------
        AgentResult
            The outcome, whose ``success`` reflects whether the verify command
            passed.
        """

run abstractmethod

run(task: AgentTask) -> AgentResult

Produce/refine task.solver_rel_path until task.verify_cmd passes.

Parameters:

Name	Type	Description	Default
task	AgentTask	The unit of work describing the workspace, the solver path to write, the restore hook, and the verify command to run.	required

Returns:

Type	Description
AgentResult	The outcome, whose success reflects whether the verify command passed.

Source code in fomtools/codegen/agent/driver.py

@abstractmethod
def run(self, task: AgentTask) -> AgentResult:
    """Produce/refine ``task.solver_rel_path`` until ``task.verify_cmd`` passes.

    Parameters
    ----------
    task : AgentTask
        The unit of work describing the workspace, the solver path to
        write, the restore hook, and the verify command to run.

    Returns
    -------
    AgentResult
        The outcome, whose ``success`` reflects whether the verify command
        passed.
    """

StubAgentDriver

Bases: AgentDriver

A deterministic AgentDriver that writes a fixed solver source.

Instead of calling an LLM, it writes a pre-supplied ("golden") solver body into the workspace and runs the real verify command, letting the full code-generation pipeline be exercised in CI without a model or network.

Methods:

Name	Description
run	Write the fixed source into the task workspace and verify it.

Source code in fomtools/codegen/agent/stub.py

class StubAgentDriver(AgentDriver):
    """A deterministic `AgentDriver` that writes a fixed solver source.

    Instead of calling an LLM, it writes a pre-supplied ("golden") solver body
    into the workspace and runs the real verify command, letting the full
    code-generation pipeline be exercised in CI without a model or network.

    Methods
    -------
    run
        Write the fixed source into the task workspace and verify it.

    """

    def __init__(
        self,
        solver_source: str | Path,
        extra_sources: dict[str, str | Callable[[str], str]] | None = None,
    ):
        """Resolve and store the solver source to write on every `run`.

        Parameters
        ----------
        solver_source : str or Path
            The golden solver body. Either C source text, or a path (``Path``
            or ``str``) to a ``.c`` file whose contents are read in. A ``str``
            is treated as a path only when it has no newlines, is under 4096
            characters, and names an existing file; otherwise it is taken as
            source text verbatim.
        extra_sources : dict, optional
            Other files to write on every run, keyed by path relative to the
            workspace: either the text to write, or a callable that maps the
            file's current text to the new one (e.g. to rewrite the generated
            linear solve). Lets tests exercise edits the pipeline must keep or
            revert.
        """
        self._extra = dict(extra_sources or {})
        if isinstance(solver_source, Path):
            self._source = solver_source.read_text()
            return
        # A str argument may be either a path to a .c file or the source itself.
        # Guard the filesystem probe: long/multi-line source confuses Path.
        self._source = solver_source
        if "\n" not in solver_source and len(solver_source) < 4096:
            try:
                path = Path(solver_source)
                if path.is_file():
                    self._source = path.read_text()
            except OSError:
                pass

    def run(self, task: AgentTask) -> AgentResult:
        """Write the fixed source into the workspace and run the verify command.

        Implements the `AgentDriver` interface. The stored source is
        written to ``task.solver_rel_path``, ``task.restore_hook`` is invoked (if present) to
        restore the deterministic files before verifying; then ``task.verify_cmd`` is run in
        ``task.workdir`` and its exit code decides success.

        Parameters
        ----------
        task : AgentTask
            The unit of work describing the workspace, the solver path to
            write, the restore hook, and the verify command to run.

        Returns
        -------
        AgentResult
            The outcome, whose ``success`` reflects whether the verify command
            passed.
        """
        target = task.workdir / task.solver_rel_path
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(self._source)
        for rel, source in self._extra.items():
            path = task.workdir / rel
            path.write_text(source(path.read_text()) if callable(source) else source)

        if task.restore_hook is not None:
            task.restore_hook()
        proc = subprocess.run(
            task.verify_cmd,
            cwd=task.workdir,
            capture_output=True,
            text=True,
        )
        log = proc.stdout + proc.stderr
        return AgentResult(success=proc.returncode == 0, log=log, turns_used=1, artifacts=[target])

__init__

__init__(solver_source: str | Path, extra_sources: dict[str, str | Callable[[str], str]] | None = None)

Parameters:

Name	Type	Description	Default
solver_source	str or Path	The golden solver body. Either C source text, or a path (Path or str) to a .c file whose contents are read in. A str is treated as a path only when it has no newlines, is under 4096 characters, and names an existing file; otherwise it is taken as source text verbatim.	required
extra_sources	dict	Other files to write on every run, keyed by path relative to the workspace: either the text to write, or a callable that maps the file's current text to the new one (e.g. to rewrite the generated linear solve). Lets tests exercise edits the pipeline must keep or revert.	None

Source code in fomtools/codegen/agent/stub.py

def __init__(
    self,
    solver_source: str | Path,
    extra_sources: dict[str, str | Callable[[str], str]] | None = None,
):
    """Resolve and store the solver source to write on every `run`.

    Parameters
    ----------
    solver_source : str or Path
        The golden solver body. Either C source text, or a path (``Path``
        or ``str``) to a ``.c`` file whose contents are read in. A ``str``
        is treated as a path only when it has no newlines, is under 4096
        characters, and names an existing file; otherwise it is taken as
        source text verbatim.
    extra_sources : dict, optional
        Other files to write on every run, keyed by path relative to the
        workspace: either the text to write, or a callable that maps the
        file's current text to the new one (e.g. to rewrite the generated
        linear solve). Lets tests exercise edits the pipeline must keep or
        revert.
    """
    self._extra = dict(extra_sources or {})
    if isinstance(solver_source, Path):
        self._source = solver_source.read_text()
        return
    # A str argument may be either a path to a .c file or the source itself.
    # Guard the filesystem probe: long/multi-line source confuses Path.
    self._source = solver_source
    if "\n" not in solver_source and len(solver_source) < 4096:
        try:
            path = Path(solver_source)
            if path.is_file():
                self._source = path.read_text()
        except OSError:
            pass

run

run(task: AgentTask) -> AgentResult

Write the fixed source into the workspace and run the verify command.

Implements the AgentDriver interface. The stored source is written to task.solver_rel_path, task.restore_hook is invoked (if present) to restore the deterministic files before verifying; then task.verify_cmd is run in task.workdir and its exit code decides success.

Parameters:

Name	Type	Description	Default
task	AgentTask	The unit of work describing the workspace, the solver path to write, the restore hook, and the verify command to run.	required

Returns:

Type	Description
AgentResult	The outcome, whose success reflects whether the verify command passed.

Source code in fomtools/codegen/agent/stub.py

def run(self, task: AgentTask) -> AgentResult:
    """Write the fixed source into the workspace and run the verify command.

    Implements the `AgentDriver` interface. The stored source is
    written to ``task.solver_rel_path``, ``task.restore_hook`` is invoked (if present) to
    restore the deterministic files before verifying; then ``task.verify_cmd`` is run in
    ``task.workdir`` and its exit code decides success.

    Parameters
    ----------
    task : AgentTask
        The unit of work describing the workspace, the solver path to
        write, the restore hook, and the verify command to run.

    Returns
    -------
    AgentResult
        The outcome, whose ``success`` reflects whether the verify command
        passed.
    """
    target = task.workdir / task.solver_rel_path
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(self._source)
    for rel, source in self._extra.items():
        path = task.workdir / rel
        path.write_text(source(path.read_text()) if callable(source) else source)

    if task.restore_hook is not None:
        task.restore_hook()
    proc = subprocess.run(
        task.verify_cmd,
        cwd=task.workdir,
        capture_output=True,
        text=True,
    )
    log = proc.stdout + proc.stderr
    return AgentResult(success=proc.returncode == 0, log=log, turns_used=1, artifacts=[target])