Diffrax

diffrax

Optional native CompiledRhs-to-Diffrax terms for smooth deterministic ODEs.

Term ownership belongs to the caller. This module creates no global solve cache, changes no JAX settings and imports neither optional package until the factory is called. Solver, save policy and adjoint selection belong to Diffrax.

DiffraxCompiledRhs

Bases: Protocol

Current compiler surface, without a compiler runtime dependency.

eval_fn property

Evaluate the flat RHS from native state and named parameters.

history_requirements property

Declare nonlocal-in-time contributions that need another driver.

meta property

Expose parameter axes, time tables and other numerical requirements.

operators property

Declare operator contributions that this scoped adapter rejects.

param_names property

Give scalar parameters; shaped parameters are declared in metadata.

pytree_eval_fn property

Evaluate native StateDict leaves without flattening, when available.

state_names property

Give the compiler's expanded flat state ordering.

template_shapes property

Give complete shapes for the native StateDict templates.

compiled_diffrax_term(compiled, *, state_layout='pytree')

Build one caller-owned Diffrax ODETerm from a current compiled RHS.

Pass a plain dict of named native JAX parameter arrays through Diffrax args. State and parameters have a common float32/float64 dtype and the compiler's exact shapes. PyTree state is a dict of base-template arrays; flat state follows compiled.state_names. No values are captured as static parameters, coerced through NumPy or placed in a global cache.

The returned object is a third-party diffrax.ODETerm. Its type remains dynamic here so importing/checking the numerical core needs no optional Diffrax installation. The caller owns the solve, output policy, adjoint, result/failure handling and smoothness assumptions. CPU operation and individual derivative transforms are qualified separately.

Returns:

Type	Description
Any	A native Diffrax ODETerm; construct it once for repeated parameter calls.

Raises:

Type	Description
ImportError	The optional Diffrax/JAX installation is absent.
ValueError	Layout metadata is inconsistent.
NotImplementedError	The RHS needs unsupported operators/history/forcing.

Source code in src/op_engine/diffrax.py

def compiled_diffrax_term(
    compiled: DiffraxCompiledRhs,
    *,
    state_layout: Literal["flat", "pytree"] = "pytree",
) -> Any:  # noqa: ANN401 - optional third-party ODETerm is the adapter boundary
    """Build one caller-owned Diffrax ODETerm from a current compiled RHS.

    Pass a plain dict of named native JAX parameter arrays through Diffrax
    ``args``. State and parameters have a common float32/float64 dtype and the
    compiler's exact shapes. PyTree state is a dict of base-template arrays;
    flat state follows ``compiled.state_names``. No values are captured as
    static parameters, coerced through NumPy or placed in a global cache.

    The returned object is a third-party ``diffrax.ODETerm``. Its type remains
    dynamic here so importing/checking the numerical core needs no optional
    Diffrax installation. The caller owns the solve, output policy, adjoint,
    result/failure handling and smoothness assumptions. CPU operation and
    individual derivative transforms are qualified separately.

    Returns:
        A native Diffrax ODETerm; construct it once for repeated parameter calls.

    Raises:
        ImportError: The optional Diffrax/JAX installation is absent.
        ValueError: Layout metadata is inconsistent.
        NotImplementedError: The RHS needs unsupported operators/history/forcing.
    """  # noqa: DOC502 - layout errors originate in the shared metadata validator
    contract = _contract(compiled, state_layout)
    try:
        diffrax = importlib.import_module("diffrax")
        jax = importlib.import_module("jax")
    except ImportError as error:
        msg = "Install op_engine[diffrax] to build a compiled Diffrax term"
        raise ImportError(msg) from error

    def evaluate(t: object, y: object, args: object) -> object:
        dtype = _state(y, contract, "state", jax)
        if not isinstance(args, dict):
            msg = "Diffrax args must be a dict of named native parameter arrays"
            raise TypeError(msg)
        if args.keys() != dict(contract.parameter_shapes).keys():
            msg = "Diffrax args keys must match compiler parameter names"
            raise ValueError(msg)
        for name, shape in contract.parameter_shapes:
            if _array(args[name], shape, f"args[{name}]", jax).dtype != dtype:
                msg = "Diffrax parameter and state dtypes must match"
                raise TypeError(msg)
        result = contract.evaluate(t, y, **args)
        if _state(result, contract, "RHS output", jax) != dtype:
            msg = "Diffrax RHS output dtype must match the state"
            raise TypeError(msg)
        return result

    return diffrax.ODETerm(evaluate)