API Reference

Solve functions

tinydiffeq.solve_ode(f, solver, t_0, t_1, x_0, *, p=None, args=None, dt_0=None, save_at=None, controller=None, max_steps=4096, project=None, has_aux=None, failure_ad_reference=None, adaptive_loop='bounded', unroll=1)

Integrate dx/dt = f(x, t, args, p) from t_0 to t_1 > t_0.

The field may be declared f(x), f(x, t), f(x, t, args), or f(x, t, args, p); x is an array or pytree with one real floating dtype, args is inert data, and p holds differentiable parameters. dt_0 is required. Fixed stepping and the default adaptive_loop="bounded" use bounded lax.scan loops with exactly max_steps attempt slots and support forward and reverse AD; adaptive_loop="forward" runs a dynamic lax.while_loop (primal, JVP, and nested forward mode only). project (an idempotent clamp) is applied at every field evaluation and accepted state. The field may return (dx, aux); saved aux follows SaveAt and participates in AD. unroll (a static int, fixed stepping only) unrolls that many steps per iteration of the integration scan — identical values, fewer/larger GPU dispatches, more compile time. Returns a :class:Solution; sol.ok reports whether t_1 was reached with every requested output valid — outputs are never poisoned.

Source code in src/tinydiffeq/ode.py

def solve_ode(
    f,
    solver,
    t_0,
    t_1,
    x_0,
    *,
    p=None,
    args=None,
    dt_0=None,
    save_at=None,
    controller=None,
    max_steps=4096,
    project=None,
    has_aux=None,
    failure_ad_reference=None,
    adaptive_loop="bounded",
    unroll=1,
):
    """Integrate ``dx/dt = f(x, t, args, p)`` from ``t_0`` to ``t_1 > t_0``.

    The field may be declared ``f(x)``, ``f(x, t)``, ``f(x, t, args)``, or
    ``f(x, t, args, p)``; ``x`` is an array or pytree with one real floating
    dtype, ``args`` is inert data, and ``p`` holds differentiable parameters.
    ``dt_0`` is required. Fixed stepping and the default
    ``adaptive_loop="bounded"`` use bounded ``lax.scan`` loops with exactly
    ``max_steps`` attempt slots and support forward and reverse AD;
    ``adaptive_loop="forward"`` runs a dynamic ``lax.while_loop`` (primal,
    JVP, and nested forward mode only). ``project`` (an idempotent clamp) is
    applied at every field evaluation and accepted state. The field may
    return ``(dx, aux)``; saved aux follows ``SaveAt`` and participates in
    AD. ``unroll`` (a static int, fixed stepping only) unrolls that many
    steps per iteration of the integration scan — identical values,
    fewer/larger GPU dispatches, more compile time. Returns a
    :class:`Solution`; ``sol.ok`` reports whether ``t_1`` was reached with
    every requested output valid — outputs are never poisoned.
    """
    if dt_0 is None:
        raise ValueError("dt_0 is required (tinydiffeq has no initial-step heuristic)")
    if adaptive_loop not in ("bounded", "forward"):
        raise ValueError(
            'adaptive_loop must be either "bounded" or "forward", got '
            f"{adaptive_loop!r}"
        )
    if save_at is None:
        save_at = SaveAt(t_1=True)
    if controller is None:
        controller = ConstantStepSize()
    if project is None:
        project = identity_project
    if controller.uses_error_estimate and not solver.has_error_estimate:
        raise ValueError(
            f"{type(controller).__name__} needs an embedded error estimate, "
            f"which {type(solver).__name__} does not provide"
        )
    if adaptive_loop == "forward" and not controller.uses_error_estimate:
        raise ValueError(
            'adaptive_loop="forward" requires an adaptive error controller'
        )
    if not isinstance(unroll, int) or isinstance(unroll, bool) or unroll < 1:
        raise ValueError("unroll must be a static int of at least 1")
    if unroll != 1 and controller.uses_error_estimate:
        raise ValueError("unroll requires fixed stepping (ConstantStepSize)")
    f = canonicalize_field(f)
    is_rodas = isinstance(solver, Rodas5P)
    is_fixed = isinstance(controller, ConstantStepSize)
    if save_at.exact and (not is_fixed or is_rodas):
        raise ValueError(
            "SaveAt exact=True requires an explicit solver with ConstantStepSize"
        )

    original_times = (t_0, t_1, dt_0)
    x_0, time_dtype = asarray_state(x_0, "x_0")
    static_uniform_horizon = _has_static_uniform_horizon(
        *original_times, max_steps, time_dtype
    )
    t_0 = jnp.asarray(t_0, time_dtype)
    t_1 = jnp.asarray(t_1, time_dtype)
    dt_0 = jnp.asarray(dt_0, time_dtype)
    time_scale = jnp.maximum(jnp.maximum(1.0, jnp.abs(t_0)), jnp.abs(t_1))
    t_eps = 4.0 * jnp.finfo(time_dtype).eps * time_scale

    def project_state(x):
        value, dtype = asarray_state(project(x), "project(x)")
        assert_same_structure(x_0, value, "project(x)")
        if dtype != time_dtype:
            raise TypeError("project(x) must preserve the state dtype")
        return value

    has_aux, aux_shape = resolve_field_aux(
        f,
        (project_state(x_0), t_0, args, p),
        jax.tree.structure(x_0),
        has_aux,
        name="has_aux",
    )

    def field_output(x, t, p_value):
        return f(project_state(x), t, args, p_value)

    def g(x, t):
        output = field_output(x, t, p)
        value, _ = split_field_output(output, has_aux)
        value, dtype = asarray_state(value, "f(x, t)")
        assert_same_structure(x_0, value, "f(x, t)")
        if dtype != time_dtype:
            raise TypeError("f(x, t) must preserve the state dtype")
        return value

    if has_aux:
        failure_ad_reference = prepare_aux_reference(failure_ad_reference, x_0, t_0, p)

        def auxiliary(inputs):
            x_value, t_value, p_value = inputs
            output = field_output(x_value, t_value, p_value)
            return split_field_output(output, True)[1]

        evaluate_aux = make_safe_evaluator(auxiliary, aux_shape)
        zero_aux = zeros_from_shape(aux_shape)

        def aux_value_and_derivative(x, t, x_dot, active):
            (value, ok), (value_dot, _) = jax.jvp(
                lambda inputs: evaluate_aux(inputs, active, failure_ad_reference),
                ((x, t, p),),
                (((x_dot, jnp.ones_like(t), zero_tangent(p))),),
            )
            return value, ok, value_dot

    else:
        evaluate_aux = None
        zero_aux = None

    need_f = not is_rodas and (
        solver.fsal or (save_at.ts is not None and not save_at.exact)
    )
    f_init = g(x_0, t_0) if need_f else zeros_like(x_0)
    track_aux = has_aux and not save_at.t_1 and not save_at.exact
    if track_aux:
        aux_init, aux_init_ok = evaluate_aux(
            (x_0, t_0, p), jnp.asarray(True), failure_ad_reference
        )
        if save_at.ts is not None and not is_rodas:
            aux_init, aux_init_ok, aux_dot_init = aux_value_and_derivative(
                x_0, t_0, f_init, aux_init_ok
            )
        else:
            aux_dot_init = zero_aux
    else:
        aux_init = None
        aux_dot_init = None
        aux_init_ok = jnp.asarray(True)
    controller_state_init = controller.init(x_0)
    flat_x_0, _ = ravel_pytree(x_0)
    identity_mass = jnp.ones_like(flat_x_0)

    def attempt_step(carry):
        (
            t,
            x,
            aux,
            aux_dot,
            dt,
            f_cur,
            done,
            failed,
            num_accepted,
            num_steps,
            controller_state,
        ) = carry
        h, proposed_t, reaches_horizon = select_step(
            t,
            t_0,
            t_1,
            dt,
            dt_0,
            num_steps,
            constant=is_fixed,
            time_tolerance=t_eps,
        )
        if is_rodas:
            x_1, err, dense, step_ok = rodas5p_step(
                g, t, x, h, identity_mass, project_state
            )
            f_1 = f_cur
        else:
            step = solver.step_fixed if is_fixed else solver.step
            x_1, f_1, err = step(g, t, x, h, f_cur if need_f else None, project_state)
            if need_f and f_1 is None:
                f_1 = g(x_1, proposed_t)
            dense = None
            step_ok = jnp.asarray(True)
        if is_rodas:
            control_err = where(
                step_ok,
                err,
                jax.tree.map(lambda value: jnp.full_like(value, jnp.inf), err),
            )
        else:
            control_err = err
        accept, dt_next, controller_state_next = controller.adapt(
            x, x_1, control_err, h, dt, solver.order, controller_state, t_1
        )
        provisional_advance = accept & step_ok & ~done & ~failed
        if track_aux:

            def accepted_auxiliary():
                if save_at.ts is None:
                    aux_candidate, aux_ok = evaluate_aux(
                        (x_1, proposed_t, p),
                        provisional_advance,
                        failure_ad_reference,
                    )
                    return aux_candidate, aux_ok, zero_aux, zero_aux
                if is_rodas:
                    left_dot, right_dot = rodas_dense_endpoint_derivatives(
                        x, x_1, dense, h
                    )
                    _, _, aux_left_dot = aux_value_and_derivative(
                        x, t, left_dot, provisional_advance
                    )
                    aux_candidate, aux_ok, aux_right_dot = aux_value_and_derivative(
                        x_1, proposed_t, right_dot, provisional_advance
                    )
                    return aux_candidate, aux_ok, aux_left_dot, aux_right_dot
                aux_candidate, aux_ok, aux_right_dot = aux_value_and_derivative(
                    x_1, proposed_t, f_1, provisional_advance
                )
                return aux_candidate, aux_ok, aux_dot, aux_right_dot

            aux_candidate, aux_ok, aux_left_dot, aux_right_dot = jax.lax.cond(
                provisional_advance,
                accepted_auxiliary,
                lambda: (aux, jnp.asarray(True), zero_aux, zero_aux),
            )
        else:
            aux_candidate = None
            aux_ok = jnp.asarray(True)
            aux_left_dot = None
            aux_right_dot = None
        advance = provisional_advance & aux_ok
        x_new = where(advance, x_1, x)
        aux_new = where(advance, aux_candidate, aux) if track_aux else None
        aux_dot_new = (
            where(advance, aux_right_dot, aux_dot)
            if track_aux and save_at.ts is not None and not is_rodas
            else aux_dot
        )
        t_new = jnp.where(advance, proposed_t, t)
        f_new = where(advance, f_1, f_cur) if need_f else f_cur
        dt_new = jnp.where(done | failed, dt, dt_next)
        controller_state_new = jax.tree.map(
            lambda old, new: jnp.where(done | failed | ~step_ok, old, new),
            controller_state,
            controller_state_next,
        )
        done_new = done | (advance & reaches_horizon)
        failed_new = failed if controller.uses_error_estimate else failed | ~step_ok
        failed_new = failed_new | (provisional_advance & ~aux_ok)
        num_new = num_accepted + advance.astype(jnp.int32)
        num_steps_new = num_steps + (~done & ~failed).astype(jnp.int32)
        carry_new = (
            t_new,
            x_new,
            aux_new,
            aux_dot_new,
            dt_new,
            f_new,
            done_new,
            failed_new,
            num_new,
            num_steps_new,
            controller_state_new,
        )
        if save_at.t_1:
            out = None
        elif save_at.steps:
            out = (t_new, x_new, aux_new, advance)
        elif is_rodas:
            out = (
                t_new,
                x_new,
                aux_new,
                dense,
                aux_left_dot,
                aux_right_dot,
                advance,
            )
        else:
            out = (t_new, x_new, aux_new, f_new, aux_dot_new, advance)
        return carry_new, out

    def skip_step(carry):
        t, x, aux, aux_dot, _, f_cur, _, _, _, _, _ = carry
        if save_at.t_1:
            out = None
        elif save_at.steps:
            out = (t, x, aux, jnp.asarray(False))
        elif is_rodas:
            out = (
                t,
                x,
                aux,
                (zeros_like(x), zeros_like(x), zeros_like(x)),
                zero_aux,
                zero_aux,
                jnp.asarray(False),
            )
        else:
            out = (t, x, aux, f_cur, aux_dot, jnp.asarray(False))
        return carry, out

    def body(carry, _):
        # The outer cond keeps a scalar predicate under vmap (see _unvmap),
        # so the frozen tail is skipped for real once every lane finishes;
        # the inner cond preserves per-lane freezing while any lane is live.
        def live_lanes(carry):
            return jax.lax.cond(carry[6] | carry[7], skip_step, attempt_step, carry)

        return jax.lax.cond(
            unvmap_all(carry[6] | carry[7]), skip_step, live_lanes, carry
        )

    def fixed_attempt_step(carry):
        t, x, f_cur, done, num_accepted = carry
        h, t_1_step, reaches_horizon = select_step(
            t,
            t_0,
            t_1,
            dt_0,
            dt_0,
            num_accepted,
            constant=True,
            time_tolerance=t_eps,
        )
        x_1, f_1, _ = solver.step_fixed(
            g, t, x, h, f_cur if need_f else None, project_state
        )
        if need_f and f_1 is None:
            f_1 = g(x_1, t_1_step)
        done_1 = reaches_horizon
        carry_1 = (t_1_step, x_1, f_1 if need_f else f_cur, done_1, num_accepted + 1)
        if save_at.t_1:
            output = None
        elif save_at.steps or save_at.exact:
            output = (t_1_step, x_1, jnp.asarray(True))
        else:
            output = (t_1_step, x_1, f_1, jnp.asarray(True))
        return carry_1, output

    def fixed_skip_step(carry):
        t, x, f_cur, _, _ = carry
        if save_at.t_1:
            output = None
        elif save_at.steps or save_at.exact:
            output = (t, x, jnp.asarray(False))
        else:
            output = (t, x, f_cur, jnp.asarray(False))
        return carry, output

    def fixed_body(carry, _):
        def live_lanes(carry):
            return jax.lax.cond(carry[3], fixed_skip_step, fixed_attempt_step, carry)

        return jax.lax.cond(unvmap_all(carry[3]), fixed_skip_step, live_lanes, carry)

    def uniform_fixed_body(carry, step_index):
        t, x, f_cur, _, num_accepted = carry
        h, t_1_step, reaches_horizon = select_step(
            t,
            t_0,
            t_1,
            dt_0,
            dt_0,
            step_index,
            constant=True,
            time_tolerance=t_eps,
        )
        x_1, f_1, _ = solver.step_fixed(
            g, t, x, h, f_cur if need_f else None, project_state
        )
        if need_f and f_1 is None:
            f_1 = g(x_1, t_1_step)
        carry_1 = (
            t_1_step,
            x_1,
            f_1 if need_f else f_cur,
            reaches_horizon,
            num_accepted + 1,
        )
        if save_at.t_1:
            output = None
        elif save_at.steps or save_at.exact:
            output = (t_1_step, x_1, jnp.asarray(True))
        else:
            output = (t_1_step, x_1, f_1, jnp.asarray(True))
        return carry_1, output

    def bounded_adaptive_scan(carry):
        chunk_size = min(ADAPTIVE_SCAN_CHUNK_SIZE, max_steps)
        num_chunks = (max_steps + chunk_size - 1) // chunk_size
        padded_steps = num_chunks * chunk_size
        valid = (jnp.arange(padded_steps) < max_steps).reshape(num_chunks, chunk_size)

        def repeat_output(output):
            if output is None:
                return None
            return jax.tree.map(
                lambda value: jnp.broadcast_to(value, (chunk_size,) + value.shape),
                output,
            )

        def run_chunk(chunk_carry, chunk_valid):
            def inner(inner_carry, is_valid):
                return jax.lax.cond(
                    is_valid,
                    body,
                    lambda value, _: skip_step(value),
                    inner_carry,
                    None,
                )

            return jax.lax.scan(inner, chunk_carry, chunk_valid, unroll=4)

        def skip_chunk(chunk_carry, chunk_valid):
            _, output = skip_step(chunk_carry)
            return chunk_carry, repeat_output(output)

        def outer(chunk_carry, chunk_valid):
            inactive = unvmap_all(chunk_carry[6] | chunk_carry[7])
            return jax.lax.cond(
                inactive, skip_chunk, run_chunk, chunk_carry, chunk_valid
            )

        final_carry, chunk_rows = jax.lax.scan(outer, carry, valid)
        if chunk_rows is None:
            return final_carry, None
        rows = jax.tree.map(
            lambda value: value.reshape((padded_steps,) + value.shape[2:])[:max_steps],
            chunk_rows,
        )
        return final_carry, rows

    carry_0 = (
        t_0,
        x_0,
        aux_init,
        aux_dot_init,
        dt_0,
        f_init,
        jnp.asarray(False),
        ~aux_init_ok,
        jnp.asarray(0, jnp.int32),
        jnp.asarray(0, jnp.int32),
        controller_state_init,
    )
    use_fast_fixed = is_fixed and not is_rodas and not track_aux
    if use_fast_fixed:
        fixed_carry_0 = (
            t_0,
            x_0,
            f_init,
            jnp.asarray(False),
            jnp.asarray(0, jnp.int32),
        )
        if static_uniform_horizon:
            step_indices = jnp.arange(max_steps, dtype=jnp.int32)
            fixed_final, rows = jax.lax.scan(
                uniform_fixed_body, fixed_carry_0, step_indices, unroll=unroll
            )
        else:
            fixed_final, rows = jax.lax.scan(
                fixed_body, fixed_carry_0, None, length=max_steps, unroll=unroll
            )
        t_final, x_final, _, done, num_accepted = fixed_final
        num_steps = num_accepted
        failed = jnp.asarray(False)
    elif controller.uses_error_estimate and adaptive_loop == "forward":
        final_carry, rows = forward_adaptive_while(
            carry_0,
            attempt_step=attempt_step,
            skip_step=skip_step,
            terminated=lambda carry: carry[6] | carry[7],
            max_steps=max_steps,
        )
        (
            t_final,
            x_final,
            _,
            _,
            _,
            _,
            done,
            failed,
            num_accepted,
            num_steps,
            _,
        ) = final_carry
    elif controller.uses_error_estimate:
        final_carry, rows = bounded_adaptive_scan(carry_0)
        (
            t_final,
            x_final,
            _,
            _,
            _,
            _,
            done,
            failed,
            num_accepted,
            num_steps,
            _,
        ) = final_carry
    else:
        final_carry, rows = jax.lax.scan(
            body, carry_0, None, length=max_steps, unroll=unroll
        )
        (
            t_final,
            x_final,
            _,
            _,
            _,
            _,
            done,
            failed,
            num_accepted,
            num_steps,
            _,
        ) = final_carry
    integration_ok = done & ~failed

    if save_at.t_1:
        if has_aux:
            aux_final, aux_ok = evaluate_aux(
                (x_final, t_final, p), jnp.asarray(True), failure_ad_reference
            )
        else:
            aux_final = None
            aux_ok = jnp.asarray(True)
        return Solution(
            ts=t_final,
            xs=x_final,
            ok=integration_ok & aux_ok,
            num_accepted=num_accepted,
            num_steps=num_steps,
            aux=aux_final,
        )

    if save_at.steps or save_at.exact:
        if use_fast_fixed:
            ts_s, xs_s, adv_s = rows
            aux_s = None
        else:
            ts_s, xs_s, aux_s, adv_s = rows
    elif is_rodas:
        (
            ts_s,
            xs_s,
            aux_s,
            dense_s,
            aux_left_dots_s,
            aux_right_dots_s,
            adv_s,
        ) = rows
    else:
        if use_fast_fixed:
            ts_s, xs_s, fs_s, adv_s = rows
            aux_s = None
            aux_dots_s = None
        else:
            ts_s, xs_s, aux_s, fs_s, aux_dots_s, adv_s = rows
    all_times = jnp.concatenate([t_0[None], ts_s])
    all_states = prepend(x_0, xs_s)
    if save_at.exact:
        query_times = jnp.asarray(save_at.ts, time_dtype)
        uniform_indices = jnp.rint((query_times - t_0) / dt_0).astype(jnp.int32)
        uniform_indices = jnp.clip(uniform_indices, 0, max_steps)
        uniform_times = all_times[uniform_indices]
        final_index = jnp.clip(num_accepted, 0, max_steps)
        final_times = jnp.broadcast_to(all_times[final_index], query_times.shape)
        use_final = jnp.abs(query_times - final_times) < jnp.abs(
            query_times - uniform_times
        )
        query_indices = jnp.where(use_final, final_index, uniform_indices)
        exact_times = all_times[query_indices]
        alignment_scale = jnp.maximum(
            jnp.maximum(1.0, jnp.abs(t_0)),
            jnp.maximum(jnp.abs(query_times), jnp.abs(exact_times)),
        )
        alignment_tolerance = jnp.minimum(
            8.0 * jnp.finfo(time_dtype).eps * alignment_scale,
            0.25 * jnp.abs(dt_0),
        )
        aligned = (
            (query_times >= t_0)
            & (query_times <= t_final)
            & (jnp.abs(query_times - exact_times) <= alignment_tolerance)
            & (query_indices <= num_accepted)
        )
        query_states = take(all_states, query_indices)
        if has_aux:

            def exact_auxiliary(x_value, t_value):
                return evaluate_aux(
                    (x_value, t_value, p),
                    jnp.asarray(True),
                    failure_ad_reference,
                )

            query_aux, query_aux_ok = jax.vmap(exact_auxiliary)(
                query_states, exact_times
            )
            aux_ok = jnp.all(query_aux_ok)
        else:
            query_aux = None
            aux_ok = jnp.asarray(True)
        return Solution(
            ts=query_times,
            xs=query_states,
            ok=integration_ok & jnp.all(aligned) & aux_ok,
            num_accepted=num_accepted,
            num_steps=num_steps,
            aux=query_aux,
        )
    all_aux = prepend(aux_init, aux_s) if has_aux else None
    raw_accepted = jnp.concatenate([jnp.ones((1,), bool), adv_s])

    if save_at.steps:
        output_size = max_steps + 1
        accepted_indices = jnp.nonzero(raw_accepted, size=output_size, fill_value=0)[0]
        compact_times = all_times[accepted_indices]
        compact_states = take(all_states, accepted_indices)
        compact_aux = take(all_aux, accepted_indices) if has_aux else None
        accepted = jnp.arange(output_size) <= num_accepted
        last_time = compact_times[num_accepted]
        last_state = take(compact_states, num_accepted)
        last_aux = take(compact_aux, num_accepted) if has_aux else None
        if save_at.fill == "inf":
            output_times = jnp.where(accepted, compact_times, jnp.inf)
        else:
            output_times = jnp.where(accepted, compact_times, last_time)
        output_states = fill_rows(compact_states, accepted, last_state, save_at.fill)
        return Solution(
            ts=output_times,
            xs=output_states,
            ok=integration_ok,
            num_accepted=num_accepted,
            num_steps=num_steps,
            accepted=accepted,
            aux=(
                fill_rows(compact_aux, accepted, last_aux, save_at.fill)
                if has_aux
                else None
            ),
        )

    query_times = jnp.asarray(save_at.ts, time_dtype)
    if is_rodas:
        query_states = rodas_interpolate(query_times, all_times, all_states, dense_s)
        query_aux = (
            hermite_interval_interpolate(
                query_times,
                all_times,
                all_aux,
                aux_left_dots_s,
                aux_right_dots_s,
            )
            if has_aux
            else None
        )
    else:
        fs_all = prepend(f_init, fs_s)
        query_states = hermite_interpolate(query_times, all_times, all_states, fs_all)
        aux_dots_all = prepend(aux_dot_init, aux_dots_s) if has_aux else None
        query_aux = (
            hermite_interpolate(query_times, all_times, all_aux, aux_dots_all)
            if has_aux
            else None
        )
    return Solution(
        ts=query_times,
        xs=query_states,
        ok=integration_ok,
        num_accepted=num_accepted,
        num_steps=num_steps,
        aux=query_aux,
    )

tinydiffeq.solve_bvp(fun, bc, t, y_0, z_0=None, *, p=None, args=None, S=None, fun_jac_ad='auto', bc_jac_ad='auto', tol=0.001, bc_tol=None, max_nodes=128, has_aux=None)

Solve dy/dt = fun(t, y, z, args, p) + S y / (t - t[0]) with two-point boundary conditions bc(y(t_a), y(t_b), z, args, p) = 0.

A faithful JAX port of :func:scipy.integrate.solve_bvp (4th-order Lobatto IIIA collocation with residual-controlled mesh refinement and a damped Newton method), with scipy's algorithm, constants, and default tolerances. fun and bc are pointwise — a scalar t and one node's state pytree — and may be declared with 2 to 5 positional arguments in the orders above. z_0 is the guess for scipy's unknown parameters (any pytree), solved jointly with y and returned as sol.z; bc must then return n + size(z) residuals as a 1-D array. p holds known differentiable parameters — the only AD input: JVP/VJP rules (composing to higher order) differentiate sol.y, sol.yp, sol.z, and sol.aux with respect to p implicitly at the solution, never through the iterations, and the guesses t, y_0, z_0 (and args, S) are differentiation-inert. args is inert pass-through data. Local Jacobians come from AD instead of scipy's finite differences; fun_jac_ad/bc_jac_ad choose "jvp", "vjp", or "auto" (forward when square or tall, reverse when strictly fat). max_nodes (static, default 128) fixes the padded output length: the mesh t starts from the given guess and grows under refinement, the returned tail repeats t[-1] and the last active rows, and sol.num_nodes counts active nodes, so hermite_interpolate(ts, sol.t, sol.y, sol.yp) evaluates exactly scipy's returned C1 cubic spline (hermite_derivative its derivative). fun may return (value, aux); aux is evaluated once at the solution over all padded nodes and participates in AD. tol is floored at 100 * eps of the working dtype (taken from y_0), silently. Failures never raise inside the solve: sol.status carries scipy's codes and a singular collocation Jacobian is reported as status 2 with the last iterate returned. The collocation system is factored once per Newton refresh by a structured orthogonal factorization of its bordered almost-block-diagonal form (O(max_nodes) instead of the dense cubic, where scipy uses sparse LU). The whole solve is compiled with lax.while_loop loops, keyed on the identity of fun and bc (reuse module-level functions rather than rebuilding closures per call); for repeated solves call it inside an outer jax.jit so the wrapper's per-call validation and dispatch trace away.

Source code in src/tinydiffeq/bvp.py

def solve_bvp(
    fun,
    bc,
    t,
    y_0,
    z_0=None,
    *,
    p=None,
    args=None,
    S=None,
    fun_jac_ad="auto",
    bc_jac_ad="auto",
    tol=1e-3,
    bc_tol=None,
    max_nodes=128,
    has_aux=None,
):
    """Solve ``dy/dt = fun(t, y, z, args, p) + S y / (t - t[0])`` with
    two-point boundary conditions ``bc(y(t_a), y(t_b), z, args, p) = 0``.

    A faithful JAX port of :func:`scipy.integrate.solve_bvp` (4th-order Lobatto
    IIIA collocation with residual-controlled mesh refinement and a damped
    Newton method), with scipy's algorithm, constants, and default tolerances.
    ``fun`` and ``bc`` are pointwise — a scalar ``t`` and one node's state
    pytree — and may be declared with 2 to 5 positional arguments in the
    orders above. ``z_0`` is the guess for scipy's unknown parameters (any
    pytree), solved jointly with ``y`` and returned as ``sol.z``; ``bc`` must
    then return ``n + size(z)`` residuals as a 1-D array. ``p`` holds known
    differentiable parameters — the only AD input: JVP/VJP rules (composing
    to higher order) differentiate ``sol.y``, ``sol.yp``, ``sol.z``, and
    ``sol.aux`` with respect to ``p`` implicitly at the solution, never
    through the iterations, and the guesses ``t``, ``y_0``, ``z_0`` (and
    ``args``, ``S``) are differentiation-inert.
    ``args`` is inert pass-through data. Local Jacobians come from AD instead
    of scipy's finite differences; ``fun_jac_ad``/``bc_jac_ad`` choose
    ``"jvp"``, ``"vjp"``, or ``"auto"`` (forward when square or tall, reverse
    when strictly fat). ``max_nodes`` (static, default 128) fixes the padded
    output length: the mesh ``t`` starts from the given guess and grows under
    refinement, the returned tail repeats ``t[-1]`` and the last active rows,
    and ``sol.num_nodes`` counts active nodes, so
    ``hermite_interpolate(ts, sol.t, sol.y, sol.yp)`` evaluates exactly
    scipy's returned C1 cubic spline (``hermite_derivative`` its derivative).
    ``fun`` may return ``(value, aux)``; aux is evaluated once at the solution
    over all padded nodes and participates in AD. ``tol`` is floored at
    ``100 * eps`` of the working dtype (taken from ``y_0``), silently.
    Failures never raise inside the solve: ``sol.status`` carries scipy's
    codes and a singular collocation Jacobian is reported as status 2 with the
    last iterate returned. The collocation system is factored once per Newton
    refresh by a structured orthogonal factorization of its bordered
    almost-block-diagonal form (``O(max_nodes)`` instead of the dense cubic,
    where scipy uses sparse LU). The whole solve is compiled with
    ``lax.while_loop`` loops, keyed on the identity of ``fun`` and ``bc``
    (reuse module-level functions rather than rebuilding closures per call);
    for repeated solves call it inside an outer ``jax.jit`` so the wrapper's
    per-call validation and dispatch trace away.
    """
    if not isinstance(max_nodes, int) or isinstance(max_nodes, bool) or max_nodes < 2:
        raise ValueError("max_nodes must be a static int of at least 2")
    for name, mode in (("fun_jac_ad", fun_jac_ad), ("bc_jac_ad", bc_jac_ad)):
        if mode not in ("auto", "jvp", "vjp"):
            raise ValueError(f'{name} must be "auto", "jvp", or "vjp"')

    y_0, dtype = asarray_state(y_0, "y_0")
    t = jnp.asarray(t, dtype)
    if t.ndim != 1:
        raise ValueError("t must be 1-dimensional")
    m_0 = t.shape[0]
    if m_0 < 2:
        raise ValueError("t must contain at least two nodes")
    if m_0 > max_nodes:
        raise ValueError(f"t has {m_0} nodes, more than max_nodes={max_nodes}")
    if not isinstance(t, jax.core.Tracer):
        if np.any(np.diff(np.asarray(t)) <= 0):
            raise ValueError("t must be strictly increasing")
    for leaf in jax.tree.leaves(y_0):
        if leaf.shape[0] != m_0:
            raise ValueError("y_0 leaves must have leading axis len(t)")

    node_template = jax.tree.map(lambda leaf: leaf[0], y_0)
    flat_node, unravel_y = ravel_pytree(node_template)
    n = flat_node.size
    Y_0 = jax.vmap(lambda node: ravel_pytree(node)[0])(y_0)

    if z_0 is None:
        k = 0
        Z_0 = jnp.zeros((0,), dtype)
        unravel_z = _unravel_empty
        z_treedef = None
        z_leaf_specs = None
    else:
        z_0, z_dtype = asarray_state(z_0, "z_0")
        if z_dtype != dtype:
            raise TypeError("z_0 must have the same dtype as y_0")
        Z_0, unravel_z = ravel_pytree(z_0)
        k = Z_0.size
        z_treedef = jax.tree.structure(z_0)
        z_leaf_specs = tuple(
            (leaf.shape, str(leaf.dtype)) for leaf in jax.tree.leaves(z_0)
        )

    if S is not None:
        S = jnp.asarray(S, dtype)
        if S.shape != (n, n):
            raise ValueError(f"S must have shape {(n, n)}, got {S.shape}")

    fun_arity = bvp_arity(
        fun, "fun", "(t, y), (t, y, z), (t, y, z, args), or (t, y, z, args, p)"
    )
    bc_arity = bvp_arity(
        bc, "bc", "(ya, yb), (ya, yb, z), (ya, yb, z, args), or (ya, yb, z, args, p)"
    )
    # Silently dropping p or args would make derivatives silently zero.
    if p is not None and fun_arity < 5 and bc_arity < 5:
        raise ValueError("p was passed but neither fun nor bc takes it")
    if args is not None and fun_arity < 4 and bc_arity < 4:
        raise ValueError("args was passed but neither fun nor bc takes it")
    if z_0 is not None and fun_arity < 3 and bc_arity < 3:
        raise ValueError("z_0 was passed but neither fun nor bc takes it")
    fun_canon = canonicalize_bvp_fun(fun, fun_arity)
    bc_canon = canonicalize_bvp_bc(bc, bc_arity)

    has_aux, _ = resolve_field_aux(
        fun_canon,
        (t[0], node_template, z_0, args, p),
        jax.tree.structure(node_template),
        has_aux,
        name="has_aux",
    )
    bc_shape = jax.eval_shape(
        lambda *operands: jnp.asarray(bc_canon(*operands)),
        node_template,
        node_template,
        z_0,
        args,
        p,
    )
    if bc_shape.shape != (n + k,):
        raise ValueError(f"bc must return a 1-D array of {n + k} residuals")

    cfg = _BVPConfig(
        fun=fun,
        bc=bc,
        fun_arity=fun_arity,
        bc_arity=bc_arity,
        max_nodes=max_nodes,
        n=n,
        k=k,
        has_aux=has_aux,
        fun_jac_ad=fun_jac_ad,
        bc_jac_ad=bc_jac_ad,
        has_singular_term=S is not None,
        y_treedef=jax.tree.structure(node_template),
        y_leaf_specs=tuple(
            (leaf.shape, str(leaf.dtype)) for leaf in jax.tree.leaves(node_template)
        ),
        z_treedef=z_treedef,
        z_leaf_specs=z_leaf_specs,
        fun_canon=fun_canon,
        bc_canon=bc_canon,
        unravel_y=unravel_y,
        unravel_z=unravel_z,
    )

    pad = max_nodes - m_0
    t_pad = jnp.concatenate([t, jnp.full((pad,), t[-1], dtype)])
    Y_pad = jnp.concatenate([Y_0, jnp.broadcast_to(Y_0[-1], (pad, n))])
    bc_tol_value = jnp.asarray(jnp.nan if bc_tol is None else bc_tol, dtype)
    t_f, Y_f, Z_f, f_f, rms, num_nodes, num_iterations, status, aux = _run_with_ad(
        cfg,
        t_pad,
        Y_pad,
        Z_0,
        jnp.asarray(m_0, jnp.int32),
        p,
        args,
        S,
        jnp.asarray(tol, dtype),
        bc_tol_value,
    )
    return BVPSolution(
        t=t_f,
        y=jax.vmap(unravel_y)(Y_f),
        yp=jax.vmap(unravel_y)(f_f),
        z=unravel_z(Z_f),
        rms_residuals=rms,
        num_nodes=num_nodes,
        num_iterations=num_iterations,
        status=status,
        ok=status == STATUS_CONVERGED,
        aux=aux,
    )

tinydiffeq.solve_semi_explicit_dae(f, g, solver, t_0, t_1, y_0, z_0, *, p=None, args=None, dt_0=None, save_at=None, controller=None, root_solver=None, has_aux=None, has_algebraic_aux=None, failure_ad_reference=None, max_steps=4096, adaptive_loop='bounded')

Integrate a semi-explicit index-1 DAE.

The system is dy/dt = f(y, z, t, args, p) with 0 = g(y, z, t, args, p) and a square nonsingular dg/dz. z_0 is a root guess: initial consistency is solved automatically and its derivative comes from the constraint. RK4 and Tsit5 restore g = 0 at every stage through :class:LMRootSolver; Rodas5P advances the block mass-matrix system with one reused LU factorization per attempted step, so its later z values satisfy the constraint to integration accuracy. Roots and their implicit derivatives are delegated to nlls-gram. f may return (dy, saved_aux); g may return (residual, algebraic_aux), in which case f takes (y, z, t, args, p, algebraic_aux). failure_ad_reference=(y, z, t, p) provides a domain-safe point for inactive vmap lanes. adaptive_loop follows :func:tinydiffeq.solve_ode. Returns a :class:DAESolution with root-solve diagnostics.

Source code in src/tinydiffeq/dae.py

def solve_semi_explicit_dae(
    f,
    g,
    solver,
    t_0,
    t_1,
    y_0,
    z_0,
    *,
    p=None,
    args=None,
    dt_0=None,
    save_at=None,
    controller=None,
    root_solver=None,
    has_aux=None,
    has_algebraic_aux=None,
    failure_ad_reference=None,
    max_steps=4096,
    adaptive_loop="bounded",
):
    """Integrate a semi-explicit index-1 DAE.

    The system is ``dy/dt = f(y, z, t, args, p)`` with
    ``0 = g(y, z, t, args, p)`` and a square nonsingular ``dg/dz``. ``z_0``
    is a root guess: initial consistency is solved automatically and its
    derivative comes from the constraint. RK4 and Tsit5 restore ``g = 0`` at
    every stage through :class:`LMRootSolver`; Rodas5P advances the block
    mass-matrix system with one reused LU factorization per attempted step,
    so its later ``z`` values satisfy the constraint to integration accuracy.
    Roots and their implicit derivatives are delegated to nlls-gram. ``f``
    may return ``(dy, saved_aux)``; ``g`` may return
    ``(residual, algebraic_aux)``, in which case ``f`` takes
    ``(y, z, t, args, p, algebraic_aux)``.
    ``failure_ad_reference=(y, z, t, p)`` provides a domain-safe point for
    inactive ``vmap`` lanes. ``adaptive_loop`` follows
    :func:`tinydiffeq.solve_ode`. Returns a :class:`DAESolution` with
    root-solve diagnostics.
    """
    if dt_0 is None:
        raise ValueError("dt_0 is required (tinydiffeq has no initial-step heuristic)")
    if adaptive_loop not in ("bounded", "forward"):
        raise ValueError(
            'adaptive_loop must be either "bounded" or "forward", got '
            f"{adaptive_loop!r}"
        )
    if save_at is None:
        save_at = SaveAt(t_1=True)
    if controller is None:
        controller = ConstantStepSize()
    if root_solver is None:
        root_solver = LMRootSolver()
    if not isinstance(root_solver, LMRootSolver):
        raise TypeError("root_solver must be an LMRootSolver")
    if not isinstance(solver, (RK4, Tsit5, Rodas5P)):
        raise TypeError("semi-explicit DAEs currently support RK4, Tsit5, and Rodas5P")
    if controller.uses_error_estimate and not solver.has_error_estimate:
        raise ValueError(
            f"{type(controller).__name__} needs an embedded error estimate, "
            f"which {type(solver).__name__} does not provide"
        )
    if adaptive_loop == "forward" and not controller.uses_error_estimate:
        raise ValueError(
            'adaptive_loop="forward" requires an adaptive error controller'
        )
    if save_at.exact:
        raise ValueError("SaveAt exact=True is only supported by solve_ode")

    y_0, time_dtype = asarray_state(y_0, "y_0")
    z_0, z_dtype = asarray_state(z_0, "z_0")
    t_0 = jnp.asarray(t_0, time_dtype)
    t_1 = jnp.asarray(t_1, time_dtype)
    dt_0 = jnp.asarray(dt_0, time_dtype)
    failure_ad_reference = _prepare_failure_ad_reference(
        failure_ad_reference, y_0, z_0, t_0, p
    )
    time_scale = jnp.maximum(jnp.maximum(1.0, jnp.abs(t_0)), jnp.abs(t_1))
    t_eps = 4.0 * jnp.finfo(time_dtype).eps * time_scale

    raw_f = f
    g_field = _canonicalize_dae_field(g, "g")
    has_algebraic_aux, algebraic_aux_shape = resolve_algebraic_aux(
        g_field,
        (y_0, z_0, t_0, args, p),
        has_algebraic_aux,
    )
    if has_algebraic_aux:
        f = _canonicalize_cached_dae_field(raw_f, "f")
        f_primals = (y_0, z_0, t_0, args, p, algebraic_aux_shape)
    else:
        f = _canonicalize_dae_field(raw_f, "f")
        f_primals = (y_0, z_0, t_0, args, p)
    has_aux, aux_shape = resolve_field_aux(
        f,
        f_primals,
        jax.tree.structure(y_0),
        has_aux,
        name="has_aux",
    )
    algebraic_solver = _get_algebraic_solver(g, root_solver, has_algebraic_aux)
    solve_root_ad, residual, algebraic_auxiliary = _make_implicit_root_solver(
        g_field,
        algebraic_solver,
        root_solver,
        z_0,
        z_dtype,
        args,
        has_algebraic_aux,
    )
    if has_algebraic_aux:
        context_evaluator = make_safe_evaluator(
            algebraic_auxiliary, algebraic_aux_shape
        )
        y_ref, z_ref, t_ref, p_ref = failure_ad_reference
        context_reference, _ = context_evaluator(
            (y_ref, z_ref, t_ref, p_ref),
            jnp.asarray(True),
            failure_ad_reference,
        )

        def evaluate_context(y, z, t, p_value, active):
            return context_evaluator((y, z, t, p_value), active, failure_ad_reference)

    else:
        evaluate_context = None

    def differential_output(y, z, t, p_value, algebraic_aux=None):
        if has_algebraic_aux:
            return f(y, z, t, args, p_value, algebraic_aux)
        return f(y, z, t, args, p_value)

    if has_aux:

        def auxiliary(inputs):
            y, z, t, p_value = inputs
            if has_algebraic_aux:
                output = g_field(y, z, t, args, p_value)
                _, context = split_algebraic_output(output, True)
            else:
                context = None
            output = differential_output(y, z, t, p_value, context)
            return split_field_output(output, True)[1]

        aux_evaluator = make_safe_evaluator(auxiliary, aux_shape)

        def evaluate_aux(y, z, t, p_value, active, failure_reference):
            return aux_evaluator((y, z, t, p_value), active, failure_reference)

    else:
        evaluate_aux = None

    def solve_root(y, t, z_guess, active):
        return solve_root_ad(y, t, z_guess, p, active, failure_ad_reference)

    def differential(y, z, t, active=None):
        if active is None:
            active = jnp.asarray(True)
        if has_algebraic_aux:
            context, context_ok = evaluate_context(y, z, t, p, active)
            context_eval = where(context_ok, context, context_reference)
        else:
            context = None
            context_ok = active
            context_eval = None
        y_ref, z_ref, t_ref, p_ref = failure_ad_reference
        y_eval = where(context_ok, y, y_ref)
        z_eval = where(context_ok, z, z_ref)
        t_eval = jnp.where(context_ok, t, t_ref)
        p_eval = where(context_ok, p, p_ref)
        output = differential_output(y_eval, z_eval, t_eval, p_eval, context_eval)
        value, _ = split_field_output(output, has_aux)
        value, dtype = asarray_state(value, "f(y, z, t)")
        assert_same_structure(y_0, value, "f(y, z, t)")
        if dtype != time_dtype:
            raise TypeError("f(y, z, t) must preserve the y dtype")
        return where(context_ok, value, zeros_like(y_0)), context_ok

    def algebraic_time_derivatives(y, z, t, y_dot, active):
        """IFT time derivatives for Hermite dense output at a root."""
        y_ref, z_ref, t_ref, p_ref = failure_ad_reference
        y_eval = where(active, y, y_ref)
        z_eval = where(active, z, z_ref)
        t_eval = jnp.where(active, t, t_ref)
        p_eval = where(active, p, p_ref)
        theta, unravel = ravel_pytree(z_eval)

        def residual_theta(theta_value):
            return jnp.ravel(residual(y_eval, unravel(theta_value), t_eval, p_eval))

        jacobian = jax.jacfwd(residual_theta)(theta)

        def residual_y_t(y_value, t_value):
            return jnp.ravel(residual(y_value, z_eval, t_value, p_eval))

        rhs = jax.jvp(
            residual_y_t,
            (y_eval, t_eval),
            (y_dot, jnp.ones_like(t)),
        )[1]
        # Under vmap, a scalar cond becomes selection and this solve runs on
        # inactive lanes. Replace a failed/singular system before solving so
        # neither the primal nor its transpose sees NaNs from that lane.
        identity = jnp.eye(theta.size, dtype=theta.dtype)
        jacobian_safe = jnp.where(active, jacobian, identity)
        rhs_safe = jnp.where(active, rhs, jnp.zeros_like(rhs))
        z_dot = unravel(jnp.linalg.solve(jacobian_safe, -rhs_safe))
        z_dot = where(active, z_dot, zeros_like(z_dot))
        if not has_aux:
            return z_dot, None

        aux_dot = jax.jvp(
            lambda y_value, z_value, t_value: evaluate_aux(
                y_value,
                z_value,
                t_value,
                p,
                active,
                failure_ad_reference,
            )[0],
            (y, z, t),
            (y_dot, z_dot, jnp.ones_like(t)),
        )[1]
        aux_dot = where(active, aux_dot, zeros_like(aux_dot))
        return z_dot, aux_dot

    (
        z_initial,
        initial_root_ok,
        initial_root_solves,
        initial_root_steps,
    ) = solve_root(y_0, t_0, z_0, jnp.asarray(True))
    if has_algebraic_aux:
        _, initial_context_ok = evaluate_context(
            y_0, z_initial, t_0, p, initial_root_ok
        )
    else:
        initial_context_ok = initial_root_ok
    if has_aux and not save_at.t_1:
        aux_initial, initial_aux_ok = evaluate_aux(
            y_0,
            z_initial,
            t_0,
            p,
            initial_context_ok,
            failure_ad_reference,
        )
        initial_ok = initial_context_ok & initial_aux_ok
    else:
        aux_initial = None
        initial_ok = initial_context_ok
    track_aux = has_aux and not save_at.t_1

    if isinstance(solver, Rodas5P):
        if has_algebraic_aux:

            def combined_output(inputs):
                y, z, t, p_value = inputs
                algebraic_output = g_field(y, z, t, args, p_value)
                residual_value, context = split_algebraic_output(algebraic_output, True)
                residual_value = _asarray_residual(residual_value)
                field_output = differential_output(y, z, t, p_value, context)
                differential_value, _ = split_field_output(field_output, has_aux)
                differential_value, dtype = asarray_state(
                    differential_value, "f(y, z, t)"
                )
                assert_same_structure(y_0, differential_value, "f(y, z, t)")
                if dtype != time_dtype:
                    raise TypeError("f(y, z, t) must preserve the y dtype")
                return differential_value, residual_value, context

            combined_shape = shape_tree(
                jax.eval_shape(combined_output, (y_0, z_0, t_0, p))
            )
            combined_evaluator = make_safe_evaluator(combined_output, combined_shape)

            def combined_values(y, z, t, active):
                (differential_value, residual_value, _), ok = combined_evaluator(
                    (y, z, t, p), active, failure_ad_reference
                )
                return differential_value, residual_value, ok

        else:

            def combined_values(y, z, t, active):
                y_ref, z_ref, t_ref, p_ref = failure_ad_reference
                y_eval = where(active, y, y_ref)
                z_eval = where(active, z, z_ref)
                t_eval = jnp.where(active, t, t_ref)
                p_eval = where(active, p, p_ref)
                residual_value = _asarray_residual(
                    g_field(y_eval, z_eval, t_eval, args, p_eval)
                )
                field_output = differential_output(y_eval, z_eval, t_eval, p_eval)
                differential_value, _ = split_field_output(field_output, has_aux)
                differential_value, dtype = asarray_state(
                    differential_value, "f(y, z, t)"
                )
                assert_same_structure(y_0, differential_value, "f(y, z, t)")
                if dtype != time_dtype:
                    raise TypeError("f(y, z, t) must preserve the y dtype")
                return differential_value, residual_value, active

        return _solve_rodas5p_dae(
            combined_values,
            evaluate_aux,
            t_0,
            t_1,
            y_0,
            z_initial,
            aux_initial,
            initial_ok,
            initial_root_solves,
            initial_root_steps,
            p,
            failure_ad_reference,
            dt_0,
            save_at,
            controller,
            max_steps,
            time_dtype,
            z_dtype,
            has_aux,
            adaptive_loop,
        )

    need_f = solver.fsal or (save_at.ts is not None)
    f_initial = jax.lax.cond(
        initial_ok,
        lambda: differential(y_0, z_initial, t_0)[0],
        lambda: zeros_like(y_0),
    )
    if save_at.ts is not None:
        z_dot_initial, aux_dot_initial = algebraic_time_derivatives(
            y_0, z_initial, t_0, f_initial, initial_ok
        )
    else:
        z_dot_initial, aux_dot_initial = None, None
    controller_state_initial = controller.init(y_0)

    def stage(y_stage, t_stage, z_guess, active, need_derivative=True):
        def evaluate():
            z_stage, root_ok, root_solves, root_steps = solve_root(
                y_stage, t_stage, z_guess, active
            )
            if need_derivative:
                k_stage, field_ok = jax.lax.cond(
                    root_ok,
                    lambda: differential(y_stage, z_stage, t_stage, root_ok),
                    lambda: (zeros_like(y_stage), jnp.asarray(False)),
                )
            else:
                k_stage = zeros_like(y_stage)
                if has_algebraic_aux:
                    _, field_ok = evaluate_context(
                        y_stage, z_stage, t_stage, p, root_ok
                    )
                else:
                    field_ok = root_ok
            return z_stage, k_stage, root_ok & field_ok, root_solves, root_steps

        def skip():
            zero = jnp.asarray(0, jnp.int32)
            return z_guess, zeros_like(y_stage), jnp.asarray(False), zero, zero

        return jax.lax.cond(active, evaluate, skip)

    def predicted_stage(
        y_stage,
        t_stage,
        z_previous,
        t_base,
        z_base,
        t_latest,
        z_latest,
        has_later_stage,
        active,
        need_derivative=True,
    ):
        z_guess = z_previous
        if root_solver.predictor == "secant":
            distinct = has_later_stage & (t_latest > t_base)
            strictly_later = distinct & (t_stage > t_latest)
            denominator = jnp.where(distinct, t_latest - t_base, 1.0)
            scale = (t_stage - t_base) / denominator
            extrapolated = jax.tree.map(
                lambda base, latest: (
                    base + jnp.asarray(scale, dtype=base.dtype) * (latest - base)
                ),
                z_base,
                z_latest,
            )
            z_guess = where(strictly_later, extrapolated, z_previous)
            z_guess = jax.tree.map(jax.lax.stop_gradient, z_guess)

        z_stage, k_stage, stage_ok, root_solves, root_steps = stage(
            y_stage, t_stage, z_guess, active, need_derivative
        )
        if root_solver.predictor == "secant":
            update_latest = stage_ok & (t_stage > t_base)
            t_latest = jnp.where(update_latest, t_stage, t_latest)
            z_latest = where(update_latest, z_stage, z_latest)
            has_later_stage = has_later_stage | update_latest
        return (
            z_stage,
            k_stage,
            stage_ok,
            root_solves,
            root_steps,
            t_latest,
            z_latest,
            has_later_stage,
        )

    def rk4_step(t, y, z, h, f_cur):
        k_1 = differential(y, z, t)[0] if f_cur is None else f_cur
        t_latest = t
        z_latest = z
        has_later_stage = jnp.asarray(False)
        (
            z_2,
            k_2,
            ok_2,
            solves_2,
            steps_2,
            t_latest,
            z_latest,
            has_later_stage,
        ) = predicted_stage(
            add_scaled(y, (0.5 * h, k_1)),
            t + 0.5 * h,
            z,
            t,
            z,
            t_latest,
            z_latest,
            has_later_stage,
            True,
        )
        (
            z_3,
            k_3,
            ok_3,
            solves_3,
            steps_3,
            t_latest,
            z_latest,
            has_later_stage,
        ) = predicted_stage(
            add_scaled(y, (0.5 * h, k_2)),
            t + 0.5 * h,
            z_2,
            t,
            z,
            t_latest,
            z_latest,
            has_later_stage,
            ok_2,
        )
        (
            z_4,
            k_4,
            ok_4,
            solves_4,
            steps_4,
            t_latest,
            z_latest,
            has_later_stage,
        ) = predicted_stage(
            add_scaled(y, (h, k_3)),
            t + h,
            z_3,
            t,
            z,
            t_latest,
            z_latest,
            has_later_stage,
            ok_2 & ok_3,
        )
        y_1 = add_scaled(
            y,
            (h / 6.0, weighted_sum((k_1, k_2, k_3, k_4), (1, 2, 2, 1))),
        )
        stage_ok = ok_2 & ok_3 & ok_4
        (
            z_1,
            f_1,
            endpoint_ok,
            endpoint_solves,
            endpoint_steps,
            _,
            _,
            _,
        ) = predicted_stage(
            y_1,
            t + h,
            z_4,
            t,
            z,
            t_latest,
            z_latest,
            has_later_stage,
            stage_ok,
            need_derivative=need_f,
        )
        root_solves = solves_2 + solves_3 + solves_4 + endpoint_solves
        root_steps = steps_2 + steps_3 + steps_4 + endpoint_steps
        return (
            y_1,
            z_1,
            f_1,
            None,
            stage_ok & endpoint_ok,
            root_solves,
            root_steps,
        )

    def tsit5_step(t, y, z, h, f_cur):
        k_1 = differential(y, z, t)[0] if f_cur is None else f_cur
        ks = [k_1]
        z_stage = z
        stages_ok = jnp.asarray(True)
        root_solves = jnp.asarray(0, jnp.int32)
        root_steps = jnp.asarray(0, jnp.int32)
        t_latest = t
        z_latest = z
        has_later_stage = jnp.asarray(False)
        rows = (
            ((A_21,), C_2),
            ((A_31, A_32), C_3),
            ((A_41, A_42, A_43), C_4),
            ((A_51, A_52, A_53, A_54), C_5),
            ((A_61, A_62, A_63, A_64, A_65), C_6),
        )
        for coefficients, stage_time in rows:
            y_stage = add_scaled(y, (h, weighted_sum(ks, coefficients)))
            (
                z_stage,
                k_stage,
                root_ok,
                stage_solves,
                stage_steps,
                t_latest,
                z_latest,
                has_later_stage,
            ) = predicted_stage(
                y_stage,
                t + stage_time * h,
                z_stage,
                t,
                z,
                t_latest,
                z_latest,
                has_later_stage,
                stages_ok,
            )
            stages_ok = stages_ok & root_ok
            root_solves = root_solves + stage_solves
            root_steps = root_steps + stage_steps
            ks.append(k_stage)
        y_1 = add_scaled(y, (h, weighted_sum(ks, (B_1, B_2, B_3, B_4, B_5, B_6))))
        (
            z_1,
            k_7,
            endpoint_ok,
            endpoint_solves,
            endpoint_steps,
            _,
            _,
            _,
        ) = predicted_stage(
            y_1,
            t + h,
            z_stage,
            t,
            z,
            t_latest,
            z_latest,
            has_later_stage,
            stages_ok,
        )
        root_solves = root_solves + endpoint_solves
        root_steps = root_steps + endpoint_steps
        ks.append(k_7)
        root_ok = stages_ok & endpoint_ok
        err = jax.tree.map(
            lambda value: h * value,
            weighted_sum(ks, (E_1, E_2, E_3, E_4, E_5, E_6, E_7)),
        )
        return y_1, z_1, k_7, err, root_ok, root_solves, root_steps

    def attempt_step(carry):
        (
            t,
            y,
            z,
            aux,
            dt,
            f_cur,
            z_dot,
            aux_dot,
            reached,
            failed,
            num_accepted,
            num_steps,
            num_root_solves,
            num_root_steps,
            controller_state,
        ) = carry
        h, proposed_t, reaches_horizon = select_step(
            t,
            t_0,
            t_1,
            dt,
            dt_0,
            num_steps,
            constant=not controller.uses_error_estimate,
            time_tolerance=t_eps,
        )
        if isinstance(solver, RK4):
            (
                y_1,
                z_1,
                f_1,
                err,
                root_ok,
                attempt_root_solves,
                attempt_root_steps,
            ) = rk4_step(t, y, z, h, f_cur if need_f else None)
        else:
            (
                y_1,
                z_1,
                f_1,
                err,
                root_ok,
                attempt_root_solves,
                attempt_root_steps,
            ) = tsit5_step(t, y, z, h, f_cur if need_f else None)

        if controller.uses_error_estimate:
            control_err = where(root_ok, err, full_like(err, jnp.inf))
        else:
            control_err = err
        controller_accept, dt_next, controller_state_next = controller.adapt(
            y, y_1, control_err, h, dt, solver.order, controller_state, t_1
        )
        accept = controller_accept & root_ok
        provisional_advance = accept & ~reached & ~failed
        if track_aux:

            def accepted_aux():
                y_safe = where(provisional_advance, y_1, y)
                z_safe = where(provisional_advance, z_1, z)
                t_safe = jnp.where(provisional_advance, proposed_t, t)
                return evaluate_aux(
                    y_safe,
                    z_safe,
                    t_safe,
                    p,
                    provisional_advance,
                    failure_ad_reference,
                )

            aux_candidate, aux_ok = jax.lax.cond(
                provisional_advance,
                accepted_aux,
                lambda: (aux, jnp.asarray(True)),
            )
        else:
            aux_candidate = None
            aux_ok = jnp.asarray(True)
        advance = provisional_advance & aux_ok
        y_new = where(advance, y_1, y)
        z_new = where(advance, z_1, z)
        t_new = jnp.where(advance, proposed_t, t)
        f_new = where(advance, f_1, f_cur) if need_f else f_cur
        if track_aux:
            aux_new = where(advance, aux_candidate, aux)
        else:
            aux_new = None
        if save_at.ts is not None:

            def accepted_derivatives():
                y_safe = where(advance, y_1, y)
                z_safe = where(advance, z_1, z)
                t_safe = jnp.where(advance, proposed_t, t)
                f_safe = where(advance, f_1, f_cur)
                return algebraic_time_derivatives(
                    y_safe, z_safe, t_safe, f_safe, advance
                )

            z_dot_new, aux_dot_new = jax.lax.cond(
                advance,
                accepted_derivatives,
                lambda: (z_dot, aux_dot),
            )
        else:
            z_dot_new, aux_dot_new = None, None
        dt_new = jnp.where(reached | failed, dt, dt_next)
        controller_state_next = jax.tree.map(
            lambda old, new: jnp.where(root_ok, new, old),
            controller_state,
            controller_state_next,
        )
        if controller.uses_error_estimate:
            failed_new = failed
        else:
            failed_new = failed | ~root_ok
        failed_new = failed_new | (provisional_advance & ~aux_ok)
        reached_new = reached | (advance & reaches_horizon)
        num_new = num_accepted + advance.astype(jnp.int32)
        num_steps_new = num_steps + (~reached & ~failed).astype(jnp.int32)
        num_root_solves_new = num_root_solves + attempt_root_solves
        num_root_steps_new = num_root_steps + attempt_root_steps
        carry_new = (
            t_new,
            y_new,
            z_new,
            aux_new,
            dt_new,
            f_new,
            z_dot_new,
            aux_dot_new,
            reached_new,
            failed_new,
            num_new,
            num_steps_new,
            num_root_solves_new,
            num_root_steps_new,
            controller_state_next,
        )
        if save_at.t_1:
            out = None
        elif save_at.steps:
            out = (t_new, y_new, z_new, aux_new, advance)
        else:
            out = (
                t_new,
                y_new,
                z_new,
                aux_new,
                f_new,
                z_dot_new,
                aux_dot_new,
                advance,
            )
        return carry_new, out

    def skip_step(carry):
        t, y, z, aux, _, f_cur, z_dot, aux_dot, _, _, _, _, _, _, _ = carry
        if save_at.t_1:
            out = None
        elif save_at.steps:
            out = (t, y, z, aux, jnp.asarray(False))
        else:
            out = (
                t,
                y,
                z,
                aux,
                f_cur,
                z_dot,
                aux_dot,
                jnp.asarray(False),
            )
        return carry, out

    def body(carry, _):
        def live_lanes(carry):
            return jax.lax.cond(carry[8] | carry[9], skip_step, attempt_step, carry)

        return jax.lax.cond(
            unvmap_all(carry[8] | carry[9]), skip_step, live_lanes, carry
        )

    carry_0 = (
        t_0,
        y_0,
        z_initial,
        aux_initial,
        dt_0,
        f_initial,
        z_dot_initial,
        aux_dot_initial,
        jnp.asarray(False),
        ~initial_ok,
        jnp.asarray(0, jnp.int32),
        jnp.asarray(0, jnp.int32),
        initial_root_solves,
        initial_root_steps,
        controller_state_initial,
    )
    if controller.uses_error_estimate and adaptive_loop == "forward":
        final_carry, rows = forward_adaptive_while(
            carry_0,
            attempt_step=attempt_step,
            skip_step=skip_step,
            terminated=lambda carry: carry[8] | carry[9],
            max_steps=max_steps,
        )
    else:
        final_carry, rows = jax.lax.scan(body, carry_0, None, length=max_steps)
    (
        t_final,
        y_final,
        z_final,
        aux_final,
        _,
        _,
        _,
        _,
        reached,
        failed,
        num_accepted,
        num_steps,
        num_root_solves,
        num_root_steps,
        _,
    ) = final_carry
    integration_ok = reached & ~failed

    if save_at.t_1:
        if has_aux:
            aux_final, aux_ok = evaluate_aux(
                y_final,
                z_final,
                t_final,
                p,
                jnp.asarray(True),
                failure_ad_reference,
            )
        else:
            aux_final = None
            aux_ok = jnp.asarray(True)
        return DAESolution(
            ts=t_final,
            ys=y_final,
            zs=z_final,
            ok=integration_ok & aux_ok,
            num_accepted=num_accepted,
            num_steps=num_steps,
            aux=aux_final,
            num_root_solves=num_root_solves,
            num_root_steps=num_root_steps,
        )

    if save_at.steps:
        ts_s, ys_s, zs_s, aux_s, adv_s = rows
    else:
        ts_s, ys_s, zs_s, aux_s, fs_s, z_dots_s, aux_dots_s, adv_s = rows
    all_times = jnp.concatenate([t_0[None], ts_s])
    all_ys = prepend(y_0, ys_s)
    all_zs = prepend(z_initial, zs_s)
    all_aux = prepend(aux_initial, aux_s) if track_aux else None
    raw_accepted = jnp.concatenate([jnp.ones((1,), bool), adv_s])

    if save_at.steps:
        output_size = max_steps + 1
        accepted_indices = jnp.nonzero(raw_accepted, size=output_size, fill_value=0)[0]
        compact_times = all_times[accepted_indices]
        compact_ys = take(all_ys, accepted_indices)
        compact_zs = take(all_zs, accepted_indices)
        compact_aux = take(all_aux, accepted_indices) if track_aux else None
        accepted = jnp.arange(output_size) <= num_accepted
        last_time = compact_times[num_accepted]
        last_y = take(compact_ys, num_accepted)
        last_z = take(compact_zs, num_accepted)
        last_aux = take(compact_aux, num_accepted) if track_aux else None

        output_times = jnp.where(
            accepted,
            compact_times,
            jnp.inf if save_at.fill == "inf" else last_time,
        )
        return DAESolution(
            ts=output_times,
            ys=fill_rows(compact_ys, accepted, last_y, save_at.fill),
            zs=fill_rows(compact_zs, accepted, last_z, save_at.fill),
            ok=integration_ok,
            num_accepted=num_accepted,
            num_steps=num_steps,
            accepted=accepted,
            aux=(
                fill_rows(compact_aux, accepted, last_aux, save_at.fill)
                if track_aux
                else None
            ),
            num_root_solves=num_root_solves,
            num_root_steps=num_root_steps,
        )

    fs_all = prepend(f_initial, fs_s)
    z_dots_all = prepend(z_dot_initial, z_dots_s)
    aux_dots_all = prepend(aux_dot_initial, aux_dots_s) if track_aux else None
    query_times = jnp.asarray(save_at.ts, time_dtype)
    values = (all_ys, all_zs, all_aux) if track_aux else (all_ys, all_zs)
    derivatives = (
        (fs_all, z_dots_all, aux_dots_all) if track_aux else (fs_all, z_dots_all)
    )
    interpolated = hermite_interpolate(query_times, all_times, values, derivatives)
    if track_aux:
        query_ys, query_zs, query_aux = interpolated
    else:
        query_ys, query_zs = interpolated
        query_aux = None
    return DAESolution(
        ts=query_times,
        ys=query_ys,
        zs=query_zs,
        ok=integration_ok,
        num_accepted=num_accepted,
        num_steps=num_steps,
        aux=query_aux,
        num_root_solves=num_root_solves,
        num_root_steps=num_root_steps,
    )

tinydiffeq.solve_semi_explicit_sdae(drift, diffusion, g, solver, t_0, t_1, y_0, z_0, *, key, n_steps, p=None, args=None, save_at=None, root_solver=None, has_aux=None, has_algebraic_aux=None, failure_ad_reference=None)

Integrate a semi-explicit index-1 Ito SDAE with diagonal noise.

The system is dy = drift(y, z, t) dt + diffusion(y, z, t) dW with 0 = g(y, z, t). solver is EulerMaruyama or SRA1, applied to the reduced SDE obtained from the locally unique root z = Z(y, t): the differential state advances on a fixed uniform grid of n_steps steps, and a root solve restores consistency at every node and at SRA1's drift stage. SRA1's strong order 1.5 requires a diffusion that depends only on time. A fixed key defines one common-random-numbers path; JVP/VJP with respect to y_0 and p are pathwise, and z_0 is a root guess with zero tangent. Aux contracts, failure_ad_reference, and failure behavior follow solve_semi_explicit_dae.

Source code in src/tinydiffeq/sdae.py

def solve_semi_explicit_sdae(
    drift,
    diffusion,
    g,
    solver,
    t_0,
    t_1,
    y_0,
    z_0,
    *,
    key,
    n_steps,
    p=None,
    args=None,
    save_at=None,
    root_solver=None,
    has_aux=None,
    has_algebraic_aux=None,
    failure_ad_reference=None,
):
    """Integrate a semi-explicit index-1 Ito SDAE with diagonal noise.

    The system is ``dy = drift(y, z, t) dt + diffusion(y, z, t) dW`` with
    ``0 = g(y, z, t)``. ``solver`` is ``EulerMaruyama`` or ``SRA1``, applied
    to the reduced SDE obtained from the locally unique root ``z = Z(y, t)``:
    the differential state advances on a fixed uniform grid of ``n_steps``
    steps, and a root solve restores consistency at every node and at SRA1's
    drift stage. SRA1's strong order 1.5 requires a diffusion that depends
    only on time. A fixed ``key`` defines one common-random-numbers path;
    JVP/VJP with respect to ``y_0`` and ``p`` are pathwise, and ``z_0`` is a
    root guess with zero tangent. Aux contracts, ``failure_ad_reference``,
    and failure behavior follow ``solve_semi_explicit_dae``.
    """
    if not isinstance(n_steps, int) or isinstance(n_steps, bool):
        raise TypeError("n_steps must be a static Python int")
    if n_steps < 1:
        raise ValueError("n_steps must be at least 1")
    if not isinstance(solver, (EulerMaruyama, SRA1)):
        raise TypeError("semi-explicit SDAEs support EulerMaruyama and SRA1")
    if save_at is None:
        save_at = SaveAt(t_1=True)
    if save_at.ts is not None:
        raise ValueError(
            "SaveAt(ts=...) is not supported for SDAEs: interpolate neither "
            "rough paths nor algebraic outputs between stochastic steps"
        )
    if root_solver is None:
        root_solver = LMRootSolver()
    if not isinstance(root_solver, LMRootSolver):
        raise TypeError("root_solver must be an LMRootSolver")

    y_0, time_dtype = asarray_state(y_0, "y_0")
    z_0, z_dtype = asarray_state(z_0, "z_0")
    t_0 = jnp.asarray(t_0, time_dtype)
    t_1 = jnp.asarray(t_1, time_dtype)
    failure_ad_reference = _prepare_failure_ad_reference(
        failure_ad_reference, y_0, z_0, t_0, p
    )
    dt = (t_1 - t_0) / n_steps
    time_grid = jnp.linspace(t_0, t_1, n_steps + 1)
    noise = solver.sample_noise(y_0, key, n_steps, dt, time_dtype)
    is_sra1 = isinstance(solver, SRA1)

    raw_drift = drift
    raw_diffusion = diffusion
    g_field = _canonicalize_dae_field(g, "g")
    has_algebraic_aux, algebraic_aux_shape = resolve_algebraic_aux(
        g_field,
        (y_0, z_0, t_0, args, p),
        has_algebraic_aux,
    )
    if has_algebraic_aux:
        drift = _canonicalize_cached_dae_field(raw_drift, "drift")
        diffusion = _canonicalize_cached_dae_field(raw_diffusion, "diffusion")
        drift_primals = (y_0, z_0, t_0, args, p, algebraic_aux_shape)
    else:
        drift = _canonicalize_dae_field(raw_drift, "drift")
        diffusion = _canonicalize_dae_field(raw_diffusion, "diffusion")
        drift_primals = (y_0, z_0, t_0, args, p)
    has_aux, aux_shape = resolve_field_aux(
        drift,
        drift_primals,
        jax.tree.structure(y_0),
        has_aux,
        name="has_aux",
    )
    algebraic_solver = _get_algebraic_solver(g, root_solver, has_algebraic_aux)
    solve_root_ad, _, algebraic_auxiliary = _make_implicit_root_solver(
        g_field,
        algebraic_solver,
        root_solver,
        z_0,
        z_dtype,
        args,
        has_algebraic_aux,
    )
    if has_algebraic_aux:
        context_evaluator = make_safe_evaluator(
            algebraic_auxiliary, algebraic_aux_shape
        )

        def evaluate_context(y, z, t, active):
            return context_evaluator((y, z, t, p), active, failure_ad_reference)

    else:
        evaluate_context = None

    def drift_output(y, z, t, p_value, context=None):
        if has_algebraic_aux:
            return drift(y, z, t, args, p_value, context)
        return drift(y, z, t, args, p_value)

    def diffusion_output(y, z, t, p_value, context=None):
        if has_algebraic_aux:
            return diffusion(y, z, t, args, p_value, context)
        return diffusion(y, z, t, args, p_value)

    if has_aux:

        def auxiliary(inputs):
            y, z, t, p_value = inputs
            if has_algebraic_aux:
                output = g_field(y, z, t, args, p_value)
                _, context = split_algebraic_output(output, True)
            else:
                context = None
            return split_field_output(drift_output(y, z, t, p_value, context), True)[1]

        aux_evaluator = make_safe_evaluator(auxiliary, aux_shape)

        def evaluate_aux(y, z, t, active):
            return aux_evaluator((y, z, t, p), active, failure_ad_reference)

    else:
        evaluate_aux = None

    def solve_root(y, t, z_guess, active):
        return solve_root_ad(y, t, z_guess, p, active, failure_ad_reference)

    def checked_value(output, name):
        value, dtype = asarray_state(output, name)
        assert_same_structure(y_0, value, name)
        if dtype != time_dtype:
            raise TypeError(f"{name} must preserve the y dtype")
        return value

    (
        z_initial,
        initial_root_ok,
        initial_root_solves,
        initial_root_steps,
    ) = solve_root(y_0, t_0, z_0, jnp.asarray(True))
    if has_algebraic_aux:
        context_initial, initial_context_ok = evaluate_context(
            y_0, z_initial, t_0, initial_root_ok
        )
    else:
        context_initial = None
        initial_context_ok = initial_root_ok
    track_aux = has_aux and save_at.steps
    if track_aux:
        aux_initial, initial_aux_ok = evaluate_aux(
            y_0, z_initial, t_0, initial_context_ok
        )
        initial_ok = initial_context_ok & initial_aux_ok
    else:
        aux_initial = None
        initial_ok = initial_context_ok

    if has_algebraic_aux:
        y_ref, z_ref, t_ref, p_ref = failure_ad_reference
        context_reference, _ = context_evaluator(
            (y_ref, z_ref, t_ref, p_ref),
            jnp.asarray(True),
            failure_ad_reference,
        )
    else:
        context_reference = None

    def attempt_step(carry, inputs):
        (
            y,
            z,
            context,
            aux,
            t,
            failed,
            num_accepted,
            num_steps,
            num_root_solves,
            num_root_steps,
        ) = carry
        t_step, t_next, noise_step = inputs
        active = ~failed
        y_ref, z_ref, t_ref, p_ref = failure_ad_reference
        y_eval = where(active, y, y_ref)
        z_eval = where(active, z, z_ref)
        t_eval = jnp.where(active, t_step, t_ref)
        p_eval = where(active, p, p_ref)
        context_eval = (
            where(active, context, context_reference) if has_algebraic_aux else None
        )
        drift_raw, _ = split_field_output(
            drift_output(y_eval, z_eval, t_eval, p_eval, context_eval), has_aux
        )
        drift_value = checked_value(drift_raw, "drift(y, z, t)")
        diffusion_value = checked_value(
            diffusion_output(y_eval, z_eval, t_eval, p_eval, context_eval),
            "diffusion(y, z, t)",
        )
        if is_sra1:
            d_w_step, d_z_step = noise_step
            chi = jax.tree.map(
                lambda w, v: 0.5 * (w + v * INV_SQRT_3), d_w_step, d_z_step
            )
            # Additive-noise contract: the diffusion may depend only on time,
            # so its endpoint evaluation reuses the node's (y, z, context).
            t_next_eval = jnp.where(active, t_next, t_ref)
            diffusion_next = checked_value(
                diffusion_output(y_eval, z_eval, t_next_eval, p_eval, context_eval),
                "diffusion(y, z, t)",
            )
            t_stage = t_step + 0.75 * dt
            y_stage = add_scaled(
                y,
                (0.75 * dt, drift_value),
                (1.5, multiply(diffusion_next, chi)),
            )
            z_stage, stage_root_ok, stage_solves, stage_steps = solve_root(
                y_stage, t_stage, z, active
            )
            if has_algebraic_aux:
                context_stage, stage_context_ok = evaluate_context(
                    y_stage, z_stage, t_stage, stage_root_ok
                )
                stage_ok = stage_root_ok & stage_context_ok
            else:
                context_stage = None
                stage_ok = stage_root_ok
            y_stage_eval = where(stage_ok, y_stage, y_ref)
            z_stage_eval = where(stage_ok, z_stage, z_ref)
            t_stage_eval = jnp.where(stage_ok, t_stage, t_ref)
            context_stage_eval = (
                where(stage_ok, context_stage, context_reference)
                if has_algebraic_aux
                else None
            )
            stage_drift_raw, _ = split_field_output(
                drift_output(
                    y_stage_eval,
                    z_stage_eval,
                    t_stage_eval,
                    p_eval,
                    context_stage_eval,
                ),
                has_aux,
            )
            stage_drift_value = checked_value(stage_drift_raw, "drift(y, z, t)")
            y_candidate = add_scaled(
                y,
                (dt / 3.0, drift_value),
                (2.0 * dt / 3.0, stage_drift_value),
                (1.0, multiply(diffusion_next, d_w_step)),
                (
                    1.0,
                    multiply(add_scaled(diffusion_value, (-1.0, diffusion_next)), chi),
                ),
            )
            z_candidate, endpoint_ok, endpoint_solves, endpoint_steps = solve_root(
                y_candidate, t_next, z_stage, stage_ok
            )
            root_ok = stage_ok & endpoint_ok
            root_solves = stage_solves + endpoint_solves
            root_steps = stage_steps + endpoint_steps
        else:
            y_candidate = add_scaled(
                y,
                (dt, drift_value),
                (1.0, multiply(diffusion_value, noise_step)),
            )
            z_candidate, root_ok, root_solves, root_steps = solve_root(
                y_candidate, t_next, z, active
            )
        if has_algebraic_aux:
            context_candidate, context_ok = evaluate_context(
                y_candidate, z_candidate, t_next, root_ok & active
            )
        else:
            context_candidate = None
            context_ok = root_ok & active
        provisional_advance = root_ok & context_ok & active
        if track_aux:

            def accepted_aux():
                y_safe = where(provisional_advance, y_candidate, y)
                z_safe = where(provisional_advance, z_candidate, z)
                t_safe = jnp.where(provisional_advance, t_next, t)
                return evaluate_aux(
                    y_safe,
                    z_safe,
                    t_safe,
                    provisional_advance,
                )

            aux_candidate, aux_ok = jax.lax.cond(
                provisional_advance,
                accepted_aux,
                lambda: (aux, jnp.asarray(True)),
            )
        else:
            aux_candidate = None
            aux_ok = jnp.asarray(True)
        advance = provisional_advance & aux_ok
        y_new = where(advance, y_candidate, y)
        z_new = where(advance, z_candidate, z)
        context_new = (
            where(advance, context_candidate, context) if has_algebraic_aux else None
        )
        t_new = jnp.where(advance, t_next, t)
        if track_aux:
            aux_new = where(advance, aux_candidate, aux)
        else:
            aux_new = None
        failed_new = (
            failed
            | ~root_ok
            | (root_ok & ~context_ok)
            | (provisional_advance & ~aux_ok)
        )
        num_new = num_accepted + advance.astype(jnp.int32)
        carry_new = (
            y_new,
            z_new,
            context_new,
            aux_new,
            t_new,
            failed_new,
            num_new,
            num_steps + jnp.asarray(1, jnp.int32),
            num_root_solves + root_solves,
            num_root_steps + root_steps,
        )
        out = (t_new, y_new, z_new, aux_new, advance) if save_at.steps else None
        return carry_new, out

    def skip_step(carry, _):
        y, z, context, aux, t, failed, num_accepted, _, _, _ = carry
        out = (t, y, z, aux, jnp.asarray(False)) if save_at.steps else None
        return carry, out

    def body(carry, inputs):
        def live_lanes(pair):
            return jax.lax.cond(
                pair[0][5],
                lambda pair: skip_step(*pair),
                lambda pair: attempt_step(*pair),
                pair,
            )

        return jax.lax.cond(
            unvmap_all(carry[5]),
            lambda pair: skip_step(*pair),
            live_lanes,
            (carry, inputs),
        )

    carry_0 = (
        y_0,
        z_initial,
        context_initial,
        aux_initial,
        t_0,
        ~initial_ok,
        jnp.asarray(0, jnp.int32),
        jnp.asarray(0, jnp.int32),
        initial_root_solves,
        initial_root_steps,
    )
    (
        (
            y_final,
            z_final,
            context_final,
            aux_final,
            t_final,
            failed,
            num_accepted,
            num_steps,
            num_root_solves,
            num_root_steps,
        ),
        rows,
    ) = jax.lax.scan(
        body,
        carry_0,
        (time_grid[:-1], time_grid[1:], noise),
    )
    ok = ~failed & (num_accepted == n_steps)
    if save_at.t_1:
        if has_aux:
            aux_final, aux_ok = evaluate_aux(
                y_final, z_final, t_final, jnp.asarray(True)
            )
        else:
            aux_final = None
            aux_ok = jnp.asarray(True)
        return DAESolution(
            ts=t_final,
            ys=y_final,
            zs=z_final,
            ok=ok & aux_ok,
            num_accepted=num_accepted,
            num_steps=num_steps,
            aux=aux_final,
            num_root_solves=num_root_solves,
            num_root_steps=num_root_steps,
        )

    ts_s, ys_s, zs_s, aux_s, advance_s = rows
    all_times = jnp.concatenate([t_0[None], ts_s])
    all_ys = prepend(y_0, ys_s)
    all_zs = prepend(z_initial, zs_s)
    all_aux = prepend(aux_initial, aux_s) if track_aux else None
    accepted = jnp.concatenate([jnp.ones((1,), bool), advance_s])
    last_time = all_times[num_accepted]
    last_y = take(all_ys, num_accepted)
    last_z = take(all_zs, num_accepted)
    last_aux = take(all_aux, num_accepted) if track_aux else None
    output_times = jnp.where(
        accepted,
        all_times,
        jnp.inf if save_at.fill == "inf" else last_time,
    )
    return DAESolution(
        ts=output_times,
        ys=fill_rows(all_ys, accepted, last_y, save_at.fill),
        zs=fill_rows(all_zs, accepted, last_z, save_at.fill),
        ok=ok,
        num_accepted=num_accepted,
        num_steps=num_steps,
        accepted=accepted,
        aux=(
            fill_rows(all_aux, accepted, last_aux, save_at.fill) if track_aux else None
        ),
        num_root_solves=num_root_solves,
        num_root_steps=num_root_steps,
    )

tinydiffeq.solve_sde(drift, diffusion, solver, t_0, t_1, x_0, *, key=None, n_steps, noise=None, p=None, args=None, save_at=None, project=None, has_aux=None, failure_ad_reference=None, unroll=1)

Integrate the Ito SDE dx = drift dt + diffusion d_w with diagonal noise on a fixed grid of n_steps uniform steps from t_0 to t_1 > t_0.

drift and diffusion follow the same signature convention as solve_ode. solver is EulerMaruyama, Milstein, or SRA1, each declaring its per-step noise through solver.sample_noise(x_0, key, n_steps, dt, dtype). Exactly one of key and noise must be provided: a fixed key presamples a fixed, reproducible noise process, differentiable with respect to x_0 and p; an explicit noise pytree (validated against the solver's spec) is additionally differentiable as data. SaveAt(ts=...) raises — interpolation is wrong for rough paths. drift may return (value, aux); diffusion is value-only. unroll (a static int) unrolls that many steps per iteration of the underlying lax.scan — identical values, fewer/larger GPU dispatches, more compile time.

Source code in src/tinydiffeq/sde.py

def solve_sde(
    drift,
    diffusion,
    solver,
    t_0,
    t_1,
    x_0,
    *,
    key=None,
    n_steps,
    noise=None,
    p=None,
    args=None,
    save_at=None,
    project=None,
    has_aux=None,
    failure_ad_reference=None,
    unroll=1,
):
    """Integrate the Ito SDE ``dx = drift dt + diffusion d_w`` with diagonal
    noise on a fixed grid of ``n_steps`` uniform steps from ``t_0`` to
    ``t_1 > t_0``.

    ``drift`` and ``diffusion`` follow the same signature convention as
    ``solve_ode``. ``solver`` is ``EulerMaruyama``, ``Milstein``, or ``SRA1``,
    each declaring its per-step noise through
    ``solver.sample_noise(x_0, key, n_steps, dt, dtype)``. Exactly one of
    ``key`` and ``noise`` must be provided: a fixed ``key`` presamples a
    fixed, reproducible noise process, differentiable with respect to ``x_0``
    and ``p``; an explicit ``noise`` pytree (validated against the solver's
    spec) is additionally differentiable as data. ``SaveAt(ts=...)`` raises —
    interpolation is wrong for rough paths. ``drift`` may return
    ``(value, aux)``; ``diffusion`` is value-only. ``unroll`` (a static int)
    unrolls that many steps per iteration of the underlying ``lax.scan`` —
    identical values, fewer/larger GPU dispatches, more compile time.
    """
    if not isinstance(n_steps, int):
        raise TypeError("n_steps must be a static Python int")
    if n_steps < 1:
        raise ValueError("n_steps must be at least 1")
    if not isinstance(unroll, int) or isinstance(unroll, bool) or unroll < 1:
        raise ValueError("unroll must be a static int of at least 1")
    if (key is None) == (noise is None):
        raise ValueError("solve_sde requires exactly one of key or noise")
    if save_at is None:
        save_at = SaveAt(t_1=True)
    if save_at.ts is not None:
        raise ValueError(
            "SaveAt(ts=...) is not supported for SDEs: Hermite interpolation "
            "is wrong for rough paths; use SaveAt(t_1=True) or SaveAt(steps=True)"
        )
    if project is None:
        project = identity_project
    drift = canonicalize_field(drift, name="drift")
    diffusion = canonicalize_field(diffusion, name="diffusion")

    x_0, time_dtype = asarray_state(x_0, "x_0")
    t_0 = jnp.asarray(t_0, time_dtype)
    t_1 = jnp.asarray(t_1, time_dtype)
    dt = (t_1 - t_0) / n_steps
    if noise is None:
        noise = solver.sample_noise(x_0, key, n_steps, dt, time_dtype)
    else:
        # sample_noise only reads shapes from x_0, so the abstract trace is
        # the solver's authoritative noise spec.
        reference = jax.eval_shape(
            lambda noise_key: solver.sample_noise(
                x_0, noise_key, n_steps, dt, time_dtype
            ),
            jax.random.key(0),
        )
        noise = jax.tree.map(jnp.asarray, noise)
        if jax.tree.structure(noise) != jax.tree.structure(reference):
            raise ValueError(
                "noise must match the pytree structure of "
                "solver.sample_noise(x_0, key, n_steps, dt, dtype)"
            )
        for leaf, ref in zip(
            jax.tree.leaves(noise), jax.tree.leaves(reference), strict=True
        ):
            if leaf.shape != ref.shape:
                raise ValueError(
                    f"noise leaf shape {leaf.shape} does not match the "
                    f"solver's expected {ref.shape}"
                )
            if leaf.dtype != ref.dtype:
                raise TypeError(
                    f"noise leaf dtype {leaf.dtype} must match the state "
                    f"dtype {ref.dtype}"
                )
    time_grid = jnp.linspace(t_0, t_1, n_steps + 1)

    def project_state(x):
        value, dtype = asarray_state(project(x), "project(x)")
        assert_same_structure(x_0, value, "project(x)")
        if dtype != time_dtype:
            raise TypeError("project(x) must preserve the state dtype")
        return value

    has_aux, aux_shape = resolve_field_aux(
        drift,
        (project_state(x_0), t_0, args, p),
        jax.tree.structure(x_0),
        has_aux,
        name="has_aux",
    )

    def drift_output(x, t, p_value):
        return drift(project_state(x), t, args, p_value)

    def g_drift(x, t):
        output = drift_output(x, t, p)
        value, _ = split_field_output(output, has_aux)
        value, dtype = asarray_state(value, "drift(x, t)")
        assert_same_structure(x_0, value, "drift(x, t)")
        if dtype != time_dtype:
            raise TypeError("drift(x, t) must preserve the state dtype")
        return value

    def g_diffusion(x, t):
        value, dtype = asarray_state(
            diffusion(project_state(x), t, args, p), "diffusion(x, t)"
        )
        assert_same_structure(x_0, value, "diffusion(x, t)")
        if dtype != time_dtype:
            raise TypeError("diffusion(x, t) must preserve the state dtype")
        return value

    def body(x, inputs):
        t, noise_step = inputs
        x_1 = solver.step(g_drift, g_diffusion, t, x, dt, noise_step, project_state)
        return x_1, x_1 if save_at.steps else None

    if save_at.t_1 or not has_aux:
        x_final, step_states = jax.lax.scan(
            body, x_0, (time_grid[:-1], noise), unroll=unroll
        )
        num_accepted = jnp.asarray(n_steps, jnp.int32)
        num_steps = num_accepted
        ok = jnp.asarray(True)

    if save_at.t_1:
        if has_aux:
            failure_ad_reference = prepare_aux_reference(
                failure_ad_reference, x_0, t_0, p
            )

            def auxiliary(inputs):
                x_value, t_value, p_value = inputs
                return split_field_output(
                    drift_output(x_value, t_value, p_value), True
                )[1]

            evaluate_aux = make_safe_evaluator(auxiliary, aux_shape)
            aux_final, aux_ok = evaluate_aux(
                (x_final, t_1, p), jnp.asarray(True), failure_ad_reference
            )
        else:
            aux_final = None
            aux_ok = jnp.asarray(True)
        return Solution(
            ts=t_1,
            xs=x_final,
            ok=ok & aux_ok,
            num_accepted=num_accepted,
            num_steps=num_steps,
            aux=aux_final,
        )

    if has_aux:
        failure_ad_reference = prepare_aux_reference(failure_ad_reference, x_0, t_0, p)

        def auxiliary(inputs):
            x_value, t_value, p_value = inputs
            return split_field_output(drift_output(x_value, t_value, p_value), True)[1]

        evaluate_aux = make_safe_evaluator(auxiliary, aux_shape)
        aux_initial, initial_ok = evaluate_aux(
            (x_0, t_0, p), jnp.asarray(True), failure_ad_reference
        )

        def aux_attempt(carry, inputs):
            x, aux, t, failed, count, num_steps = carry
            t_step, t_next, noise_step = inputs
            x_candidate = solver.step(
                g_drift,
                g_diffusion,
                t_step,
                x,
                dt,
                noise_step,
                project_state,
            )
            aux_candidate, aux_ok = evaluate_aux(
                (x_candidate, t_next, p),
                ~failed,
                failure_ad_reference,
            )
            advance = ~failed & aux_ok
            x_new = where(advance, x_candidate, x)
            aux_new = where(advance, aux_candidate, aux)
            t_new = jnp.where(advance, t_next, t)
            failed_new = failed | ~aux_ok
            count_new = count + advance.astype(jnp.int32)
            num_steps_new = num_steps + jnp.asarray(1, jnp.int32)
            return (
                x_new,
                aux_new,
                t_new,
                failed_new,
                count_new,
                num_steps_new,
            ), (t_new, x_new, aux_new, advance)

        def aux_skip(carry, inputs):
            x, aux, t, failed, count, num_steps = carry
            return carry, (t, x, aux, jnp.asarray(False))

        def aux_body(carry, inputs):
            # Scalar-predicate outer cond under vmap (see _unvmap): once every
            # lane has failed, the frozen tail skips for real.
            def live_lanes(pair):
                return jax.lax.cond(
                    pair[0][3],
                    lambda pair: aux_skip(*pair),
                    lambda pair: aux_attempt(*pair),
                    pair,
                )

            return jax.lax.cond(
                unvmap_all(carry[3]),
                lambda pair: aux_skip(*pair),
                live_lanes,
                (carry, inputs),
            )

        carry_0 = (
            x_0,
            aux_initial,
            t_0,
            ~initial_ok,
            jnp.asarray(0, jnp.int32),
            jnp.asarray(0, jnp.int32),
        )
        (
            (
                x_final,
                aux_final,
                t_final,
                failed,
                num_accepted,
                num_steps,
            ),
            rows,
        ) = jax.lax.scan(
            aux_body, carry_0, (time_grid[:-1], time_grid[1:], noise), unroll=unroll
        )
        ts_s, xs_s, aux_s, advance_s = rows
        all_times = jnp.concatenate([t_0[None], ts_s])
        all_states = prepend(x_0, xs_s)
        all_aux = prepend(aux_initial, aux_s)
        accepted = jnp.concatenate([jnp.ones((1,), bool), advance_s])
        last_time = all_times[num_accepted]
        last_state = take(all_states, num_accepted)
        last_aux = take(all_aux, num_accepted)
        output_times = jnp.where(
            accepted,
            all_times,
            jnp.inf if save_at.fill == "inf" else last_time,
        )
        return Solution(
            ts=output_times,
            xs=fill_rows(all_states, accepted, last_state, save_at.fill),
            ok=~failed & (num_accepted == n_steps),
            num_accepted=num_accepted,
            num_steps=num_steps,
            accepted=accepted,
            aux=fill_rows(all_aux, accepted, last_aux, save_at.fill),
        )

    all_states = prepend(x_0, step_states)
    accepted = jnp.ones((n_steps + 1,), bool)
    return Solution(
        ts=time_grid,
        xs=all_states,
        ok=ok,
        num_accepted=num_accepted,
        num_steps=num_steps,
        accepted=accepted,
    )

tinydiffeq.solve_linear_ode(operator, method, t_0, t_1, x_0, *, save_at=None)

Solve dx/dt = A(x) for a fixed homogeneous linear operator.

operator is either a square matrix using the column convention A @ x or a callable mapping the state pytree to an identically structured pytree. DenseExponential materializes a callable operator before a dense matrix exponential; the Krylov methods evaluate only operator actions. Endpoint output is the default, and SaveAt(ts=...) evaluates independent exponential actions at the requested times. JVPs and VJPs flow through the initial state and differentiable operator arrays. The operator must be autonomous, homogeneous, and linear.

Source code in src/tinydiffeq/exponential.py

def solve_linear_ode(operator, method, t_0, t_1, x_0, *, save_at=None):
    """Solve ``dx/dt = A(x)`` for a fixed homogeneous linear operator.

    ``operator`` is either a square matrix using the column convention
    ``A @ x`` or a callable mapping the state pytree to an identically
    structured pytree. ``DenseExponential`` materializes a callable operator
    before a dense matrix exponential; the Krylov methods evaluate only
    operator actions. Endpoint output is the default, and ``SaveAt(ts=...)``
    evaluates independent exponential actions at the requested times. JVPs
    and VJPs flow through the initial state and differentiable operator
    arrays. The operator must be autonomous, homogeneous, and linear.
    """
    if not isinstance(method, _EXPONENTIAL_METHODS):
        raise TypeError(
            "method must be DenseExponential, KrylovExponential, or "
            "AdaptiveKrylovExponential"
        )
    if save_at is None:
        save_at = SaveAt(t_1=True)
    if save_at.exact:
        raise ValueError("SaveAt exact=True is only supported by solve_ode")
    if save_at.steps:
        raise ValueError("linear exponential solves require endpoint or SaveAt.ts")

    (
        x_0,
        dtype,
        flat_initial,
        unravel,
        flat_action,
        dense_operator,
    ) = _prepare_linear_operator(operator, x_0)
    t_0 = jnp.asarray(t_0, dtype)
    t_1 = jnp.asarray(t_1, dtype)

    if isinstance(method, DenseExponential) and dense_operator is None:
        dense_operator = jax.jacfwd(flat_action)(jnp.zeros_like(flat_initial))

    def evaluate(time):
        elapsed = time - t_0
        if isinstance(method, DenseExponential):
            value = _dense_exponential_action(dense_operator, flat_initial, elapsed)
            count = jnp.where(elapsed > 0, 1, 0).astype(jnp.int32)
            return value, jnp.asarray(True), count, count
        return _propagate_krylov(flat_action, flat_initial, elapsed, method)

    if save_at.t_1:
        times = t_1
        flat_states, method_ok, num_accepted, num_steps = evaluate(t_1)
        states = unravel(flat_states)
        times_ok = t_1 >= t_0
    else:
        times = jnp.asarray(save_at.ts, dtype)
        if times.ndim != 1:
            raise TypeError("SaveAt.ts must be one-dimensional")
        times_ok = jnp.all((times >= t_0) & (times <= t_1))
        flat_states, method_ok, accepted_counts, step_counts = jax.vmap(evaluate)(times)
        num_accepted = jnp.max(accepted_counts, initial=jnp.asarray(0, jnp.int32))
        num_steps = jnp.max(step_counts, initial=jnp.asarray(0, jnp.int32))
        states = jax.vmap(unravel)(flat_states)
    finite = jnp.all(jnp.isfinite(flat_states))
    return Solution(
        ts=times,
        xs=states,
        ok=times_ok & jnp.all(method_ok) & finite,
        num_accepted=num_accepted,
        num_steps=num_steps,
    )

tinydiffeq.jvp_linear_ode(operator, method, t_0, t_1, x_0, x_0_tangent, *, batched=False)

Return a terminal linear solve and hand-coded initial-state JVP.

With batched=True, every tangent leaf has a leading direction axis. Dense mode forms one exponential and applies it to every direction. Matrix-free Krylov mode vectorizes independent exponential actions. The operator is fixed: use ordinary jax.jvp when differentiating operator entries or arrays captured by a callable.

Source code in src/tinydiffeq/exponential.py

def jvp_linear_ode(
    operator,
    method,
    t_0,
    t_1,
    x_0,
    x_0_tangent,
    *,
    batched=False,
):
    """Return a terminal linear solve and hand-coded initial-state JVP.

    With ``batched=True``, every tangent leaf has a leading direction axis.
    Dense mode forms one exponential and applies it to every direction.
    Matrix-free Krylov mode vectorizes independent exponential actions. The
    operator is fixed: use ordinary ``jax.jvp`` when differentiating operator
    entries or arrays captured by a callable.
    """
    if not isinstance(method, _EXPONENTIAL_METHODS):
        raise TypeError(
            "method must be DenseExponential, KrylovExponential, or "
            "AdaptiveKrylovExponential"
        )
    (
        x_0,
        dtype,
        flat_initial,
        unravel,
        flat_action,
        dense_operator,
    ) = _prepare_linear_operator(operator, x_0)
    tangent, restore_tangent, tangent_dtype = _flatten_directions(
        x_0_tangent,
        x_0,
        unravel,
        batched=batched,
        name="x_0_tangent",
    )
    if tangent_dtype != dtype:
        raise TypeError("x_0_tangent must have the same dtype as x_0")
    if isinstance(method, DenseExponential) and dense_operator is None:
        dense_operator = jax.jacfwd(flat_action)(jnp.zeros_like(flat_initial))
    t_0 = jnp.asarray(t_0, dtype)
    t_1 = jnp.asarray(t_1, dtype)
    elapsed = t_1 - t_0
    flat_value, primal_ok, exponential, num_accepted, num_steps = _terminal_value(
        method, dense_operator, flat_action, flat_initial, elapsed
    )
    if isinstance(method, DenseExponential):
        flat_tangent = tangent @ exponential.T if batched else exponential @ tangent
        tangent_ok = jnp.asarray(True)
    elif batched:
        flat_tangent, tangent_ok, _, _ = jax.vmap(
            lambda direction: _propagate_krylov(flat_action, direction, elapsed, method)
        )(tangent)
    else:
        flat_tangent, tangent_ok, _, _ = _propagate_krylov(
            flat_action, tangent, elapsed, method
        )
    finite = jnp.all(jnp.isfinite(flat_value)) & jnp.all(jnp.isfinite(flat_tangent))
    solution = Solution(
        ts=t_1,
        xs=unravel(flat_value),
        ok=(t_1 >= t_0) & primal_ok & jnp.all(tangent_ok) & finite,
        num_accepted=num_accepted,
        num_steps=num_steps,
    )
    return solution, restore_tangent(flat_tangent)

tinydiffeq.vjp_linear_ode(operator, method, t_0, t_1, x_0, cotangent, *, batched=False)

Return a terminal linear solve and hand-coded initial-state VJP.

The pullback is another exponential action with the transposed operator. batched=True accepts multiple terminal cotangents on a leading axis. Dense mode reuses the primal exponential for every cotangent. Callable transpose actions are generated with jax.linear_transpose and therefore require the declared operator to be linear in its state argument.

Source code in src/tinydiffeq/exponential.py

def vjp_linear_ode(
    operator,
    method,
    t_0,
    t_1,
    x_0,
    cotangent,
    *,
    batched=False,
):
    """Return a terminal linear solve and hand-coded initial-state VJP.

    The pullback is another exponential action with the transposed operator.
    ``batched=True`` accepts multiple terminal cotangents on a leading axis.
    Dense mode reuses the primal exponential for every cotangent. Callable
    transpose actions are generated with ``jax.linear_transpose`` and therefore
    require the declared operator to be linear in its state argument.
    """
    if not isinstance(method, _EXPONENTIAL_METHODS):
        raise TypeError(
            "method must be DenseExponential, KrylovExponential, or "
            "AdaptiveKrylovExponential"
        )
    (
        x_0,
        dtype,
        flat_initial,
        unravel,
        flat_action,
        dense_operator,
    ) = _prepare_linear_operator(operator, x_0)
    flat_cotangent, restore_cotangent, cotangent_dtype = _flatten_directions(
        cotangent,
        x_0,
        unravel,
        batched=batched,
        name="cotangent",
    )
    if cotangent_dtype != dtype:
        raise TypeError("cotangent must have the same dtype as x_0")
    if isinstance(method, DenseExponential) and dense_operator is None:
        dense_operator = jax.jacfwd(flat_action)(jnp.zeros_like(flat_initial))
    t_0 = jnp.asarray(t_0, dtype)
    t_1 = jnp.asarray(t_1, dtype)
    elapsed = t_1 - t_0
    flat_value, primal_ok, exponential, num_accepted, num_steps = _terminal_value(
        method, dense_operator, flat_action, flat_initial, elapsed
    )
    if isinstance(method, DenseExponential):
        flat_gradient = (
            flat_cotangent @ exponential if batched else exponential.T @ flat_cotangent
        )
        gradient_ok = jnp.asarray(True)
    else:
        zero = jnp.zeros_like(flat_initial)

        def transpose_action(vector):
            return jax.linear_transpose(flat_action, zero)(vector)[0]

        if batched:
            flat_gradient, gradient_ok, _, _ = jax.vmap(
                lambda direction: _propagate_krylov(
                    transpose_action, direction, elapsed, method
                )
            )(flat_cotangent)
        else:
            flat_gradient, gradient_ok, _, _ = _propagate_krylov(
                transpose_action, flat_cotangent, elapsed, method
            )
    finite = jnp.all(jnp.isfinite(flat_value)) & jnp.all(jnp.isfinite(flat_gradient))
    solution = Solution(
        ts=t_1,
        xs=unravel(flat_value),
        ok=(t_1 >= t_0) & primal_ok & jnp.all(gradient_ok) & finite,
        num_accepted=num_accepted,
        num_steps=num_steps,
    )
    return solution, restore_cotangent(flat_gradient)

tinydiffeq.simulate_markov_chain(chain, state_0, *, key, num_steps, method=None, save_at=None)

Simulate a primal finite-state homogeneous discrete-time Markov chain.

state_0 is one scalar integer index and key one JAX key. Use jax.vmap over initial states and independently split keys for ensembles. num_steps is static. SequentialMarkov is the CPU-oriented default; AssociativeMarkov composes sampled state maps in parallel and returns the identical path for the same key. SaveAt selects the endpoint, all steps, or a one-dimensional array of integer step indices.

Source code in src/tinydiffeq/markov.py

def simulate_markov_chain(
    chain,
    state_0,
    *,
    key,
    num_steps,
    method=None,
    save_at=None,
):
    """Simulate a primal finite-state homogeneous discrete-time Markov chain.

    ``state_0`` is one scalar integer index and ``key`` one JAX key. Use
    ``jax.vmap`` over initial states and independently split keys for ensembles.
    ``num_steps`` is static. ``SequentialMarkov`` is the CPU-oriented default;
    ``AssociativeMarkov`` composes sampled state maps in parallel and returns the
    identical path for the same key. ``SaveAt`` selects the endpoint, all steps,
    or a one-dimensional array of integer step indices.
    """
    if not isinstance(chain, DiscreteMarkovChain):
        raise TypeError("chain must be a DiscreteMarkovChain")
    if method is None:
        method = SequentialMarkov()
    state_0, initial_ok, save_at = _validate_simulation_inputs(
        chain, state_0, num_steps, "num_steps", save_at
    )
    uniforms = _unit_uniform(key, (num_steps,), chain.transition_matrix.dtype)
    step_states = _simulate_discrete_states(chain, state_0, uniforms, method)
    all_states = prepend(state_0, step_states)
    count = jnp.asarray(num_steps, jnp.int32)
    if save_at.t_1:
        return Solution(
            ts=count,
            xs=step_states[-1],
            ok=initial_ok,
            num_accepted=count,
            num_steps=count,
        )
    if save_at.steps:
        return Solution(
            ts=jnp.arange(num_steps + 1, dtype=jnp.int32),
            xs=all_states,
            ok=initial_ok,
            num_accepted=count,
            num_steps=count,
            accepted=jnp.ones((num_steps + 1,), dtype=bool),
        )
    query_steps = jnp.asarray(save_at.ts)
    if query_steps.ndim != 1 or not jnp.issubdtype(query_steps.dtype, jnp.integer):
        raise TypeError("discrete SaveAt.ts must be a one-dimensional integer array")
    queries_ok = jnp.all((query_steps >= 0) & (query_steps <= num_steps))
    safe_queries = jnp.clip(query_steps, 0, num_steps).astype(jnp.int32)
    return Solution(
        ts=query_steps,
        xs=all_states[safe_queries],
        ok=initial_ok & queries_ok,
        num_accepted=count,
        num_steps=count,
    )

tinydiffeq.simulate_continuous_time_markov_chain(chain, t_0, t_1, state_0, *, key, max_jumps, method=None, save_at=None)

Simulate a primal finite-state CTMC with Gillespie's direct recurrence.

max_jumps is a static bound. Endpoint output is the default; SaveAt(steps=True) returns the padded event path and SaveAt(ts=...) evaluates its right-continuous piecewise-constant state. sol.ok is false if the jump budget does not cover t_1. Associative execution composes state and holding-time maps; states agree with sequential execution for the same key, while event times can differ by floating-point reassociation.

Source code in src/tinydiffeq/markov.py

def simulate_continuous_time_markov_chain(
    chain,
    t_0,
    t_1,
    state_0,
    *,
    key,
    max_jumps,
    method=None,
    save_at=None,
):
    """Simulate a primal finite-state CTMC with Gillespie's direct recurrence.

    ``max_jumps`` is a static bound. Endpoint output is the default;
    ``SaveAt(steps=True)`` returns the padded event path and ``SaveAt(ts=...)``
    evaluates its right-continuous piecewise-constant state. ``sol.ok`` is false
    if the jump budget does not cover ``t_1``. Associative execution composes
    state and holding-time maps; states agree with sequential execution for the
    same key, while event times can differ by floating-point reassociation.
    """
    if not isinstance(chain, ContinuousTimeMarkovChain):
        raise TypeError("chain must be a ContinuousTimeMarkovChain")
    if method is None:
        method = SequentialMarkov()
    state_0, initial_ok, save_at = _validate_simulation_inputs(
        chain, state_0, max_jumps, "max_jumps", save_at
    )
    dtype = chain.generator.dtype
    t_0 = jnp.asarray(t_0, dtype)
    t_1 = jnp.asarray(t_1, dtype)
    time_ok = t_1 >= t_0
    exponential_key, transition_key = jax.random.split(key)
    exponentials = _unit_exponential(exponential_key, (max_jumps,), dtype)
    uniforms = _unit_uniform(transition_key, (max_jumps,), dtype)
    states_after, elapsed_event_times = _simulate_continuous_events(
        chain, state_0, exponentials, uniforms, method
    )
    event_times = t_0 + elapsed_event_times
    num_jumps = jnp.sum(event_times <= t_1, dtype=jnp.int32)
    # One event beyond the horizon is needed to establish endpoint coverage;
    # a starved event budget has already attempted every available event.
    num_steps = jnp.minimum(num_jumps + 1, jnp.asarray(max_jumps, jnp.int32))
    all_states = prepend(state_0, states_after)
    final_state = all_states[num_jumps]
    covered = event_times[-1] >= t_1
    integration_ok = initial_ok & time_ok & covered

    if save_at.t_1:
        reached_time = jnp.where(integration_ok, t_1, event_times[-1])
        return Solution(
            ts=reached_time,
            xs=final_state,
            ok=integration_ok,
            num_accepted=num_jumps,
            num_steps=num_steps,
        )
    if save_at.steps:
        accepted = jnp.arange(max_jumps + 1) <= num_jumps
        raw_times = jnp.concatenate([t_0[None], event_times])
        output_times = jnp.where(accepted, raw_times, t_1)
        output_states = jnp.where(accepted, all_states, final_state)
        return Solution(
            ts=output_times,
            xs=output_states,
            ok=integration_ok,
            num_accepted=num_jumps,
            num_steps=num_steps,
            accepted=accepted,
        )
    query_times = jnp.asarray(save_at.ts, dtype)
    if query_times.ndim != 1:
        raise TypeError("continuous-time SaveAt.ts must be one-dimensional")
    queries_ok = jnp.all((query_times >= t_0) & (query_times <= t_1))
    event_counts = jnp.searchsorted(
        event_times, query_times, side="right", method="compare_all"
    )
    event_counts = jnp.minimum(event_counts, num_jumps)
    return Solution(
        ts=query_times,
        xs=all_states[event_counts],
        ok=integration_ok & queries_ok,
        num_accepted=num_jumps,
        num_steps=num_steps,
    )

tinydiffeq.forecast_markov_chain(chain, distribution_0, *, num_steps, method=None, save_at=None)

Forecast a fixed DTMC probability mass function.

Endpoint output defaults to binary matrix powering. Multi-row output defaults to chronological matrix-vector scan; AssociativeMarkov instead composes prefix transition matrices and can expose useful GPU parallelism for small state spaces. The deterministic forecast supports JVP/VJP with respect to distribution_0; the prepared chain is treated as fixed.

Source code in src/tinydiffeq/markov.py

def forecast_markov_chain(
    chain,
    distribution_0,
    *,
    num_steps,
    method=None,
    save_at=None,
):
    """Forecast a fixed DTMC probability mass function.

    Endpoint output defaults to binary matrix powering. Multi-row output defaults
    to chronological matrix-vector scan; ``AssociativeMarkov`` instead composes
    prefix transition matrices and can expose useful GPU parallelism for small
    state spaces. The deterministic forecast supports JVP/VJP with respect to
    ``distribution_0``; the prepared chain is treated as fixed.
    """
    if not isinstance(chain, DiscreteMarkovChain):
        raise TypeError("chain must be a DiscreteMarkovChain")
    if not isinstance(num_steps, int) or isinstance(num_steps, bool):
        raise TypeError("num_steps must be a static Python int")
    if num_steps < 0:
        raise ValueError("num_steps must be nonnegative")
    if save_at is None:
        save_at = SaveAt(t_1=True)
    _reject_exact_save_at(save_at)
    distribution_0, initial_ok = _prepare_distribution(chain, distribution_0)

    if save_at.t_1:
        if method is None or isinstance(method, MatrixPowerMarkov):
            probabilities = distribution_0 @ jnp.linalg.matrix_power(
                chain.transition_matrix, num_steps
            )
        elif isinstance(method, SequentialMarkov):
            probabilities = jax.lax.fori_loop(
                0,
                num_steps,
                lambda _, distribution: distribution @ chain.transition_matrix,
                distribution_0,
            )
        elif isinstance(method, AssociativeMarkov):
            if num_steps == 0:
                probabilities = distribution_0
            else:
                probabilities = _discrete_distribution_steps(
                    chain, distribution_0, num_steps, method
                )[-1]
        else:
            raise TypeError(
                "method must be MatrixPowerMarkov, SequentialMarkov, or "
                "AssociativeMarkov"
            )
        return MarkovDistribution(
            ts=jnp.asarray(num_steps, jnp.int32),
            probabilities=probabilities,
            ok=initial_ok & _flat_distribution_is_valid(probabilities),
        )

    if isinstance(method, MatrixPowerMarkov):
        raise ValueError("MatrixPowerMarkov supports endpoint output only")
    if method is None:
        method = SequentialMarkov()
    if num_steps == 0:
        all_probabilities = distribution_0[None]
    else:
        step_probabilities = _discrete_distribution_steps(
            chain, distribution_0, num_steps, method
        )
        all_probabilities = prepend(distribution_0, step_probabilities)
    if save_at.steps:
        times = jnp.arange(num_steps + 1, dtype=jnp.int32)
        probabilities = all_probabilities
    else:
        query_steps = jnp.asarray(save_at.ts)
        if query_steps.ndim != 1 or not jnp.issubdtype(query_steps.dtype, jnp.integer):
            raise TypeError(
                "discrete SaveAt.ts must be a one-dimensional integer array"
            )
        queries_ok = jnp.all((query_steps >= 0) & (query_steps <= num_steps))
        safe_queries = jnp.clip(query_steps, 0, num_steps).astype(jnp.int32)
        times = query_steps
        probabilities = all_probabilities[safe_queries]
        initial_ok = initial_ok & queries_ok
    return MarkovDistribution(
        ts=times,
        probabilities=probabilities,
        ok=initial_ok & _flat_distribution_is_valid(probabilities),
    )

tinydiffeq.forecast_continuous_time_markov_chain(chain, t_0, t_1, distribution_0, *, method=None, save_at=None)

Forecast a fixed CTMC probability mass with exponential actions.

For one endpoint this evaluates distribution_0 @ exp((t_1-t_0) Q). DenseExponential forms the dense exponential. KrylovExponential applies a static Arnoldi approximation; AdaptiveKrylovExponential adapts its internal time slices. Both support MatrixFreeContinuousTimeMarkovChain probability pytrees. Requested times are independent exponential actions vectorized over the query axis. JVP/VJP with respect to distribution_0 are supported.

Source code in src/tinydiffeq/markov.py

def forecast_continuous_time_markov_chain(
    chain,
    t_0,
    t_1,
    distribution_0,
    *,
    method=None,
    save_at=None,
):
    """Forecast a fixed CTMC probability mass with exponential actions.

    For one endpoint this evaluates ``distribution_0 @ exp((t_1-t_0) Q)``.
    ``DenseExponential`` forms the dense exponential.
    ``KrylovExponential`` applies a static Arnoldi approximation;
    ``AdaptiveKrylovExponential`` adapts its internal time slices. Both support
    ``MatrixFreeContinuousTimeMarkovChain`` probability pytrees.
    Requested times are independent exponential actions vectorized over the query
    axis. JVP/VJP with respect to ``distribution_0`` are supported.
    """
    is_dense = isinstance(chain, ContinuousTimeMarkovChain)
    is_matrix_free = isinstance(chain, MatrixFreeContinuousTimeMarkovChain)
    if not is_dense and not is_matrix_free:
        raise TypeError(
            "chain must be ContinuousTimeMarkovChain or "
            "MatrixFreeContinuousTimeMarkovChain"
        )
    if method is None:
        method = DenseExponential() if is_dense else KrylovExponential()
    if not isinstance(
        method, (DenseExponential, KrylovExponential, AdaptiveKrylovExponential)
    ):
        raise TypeError(
            "CTMC distribution forecasts require DenseExponential, "
            "KrylovExponential, or AdaptiveKrylovExponential"
        )
    if save_at is None:
        save_at = SaveAt(t_1=True)
    _reject_exact_save_at(save_at)
    if save_at.steps:
        raise ValueError("CTMC distribution forecasts require endpoint or SaveAt.ts")
    if is_dense:
        distribution_0, initial_ok = _prepare_distribution(chain, distribution_0)
        dtype = chain.generator.dtype
    else:
        distribution_0, initial_ok, dtype = _prepare_pytree_distribution(distribution_0)
    if is_dense:
        operator = chain.generator.T
    else:
        operator = chain.forward_generator
    linear_solution = solve_linear_ode(
        operator,
        method,
        t_0,
        t_1,
        distribution_0,
        save_at=save_at,
    )
    probabilities = linear_solution.xs
    return MarkovDistribution(
        ts=linear_solution.ts,
        probabilities=probabilities,
        ok=(
            initial_ok
            & linear_solution.ok
            & _pytree_distribution_is_valid(probabilities, batched=not save_at.t_1)
        ),
    )

Solvers

tinydiffeq.Euler dataclass

Explicit Euler. Fixed-step only: no embedded error estimate.

Source code in src/tinydiffeq/solvers.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class Euler:
    """Explicit Euler. Fixed-step only: no embedded error estimate."""

    order = 1
    fsal = False
    has_error_estimate = False

    def step(self, g, t, x, dt, f_0, project):
        k_1 = g(x, t) if f_0 is None else f_0
        x_1 = project(add_scaled(x, (dt, k_1)))
        return x_1, None, None

    def step_fixed(self, g, t, x, dt, f_0, project):
        return self.step(g, t, x, dt, f_0, project)

tinydiffeq.RK4 dataclass

Classic fourth-order Runge-Kutta. Fixed-step only: no error estimate.

Source code in src/tinydiffeq/solvers.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class RK4:
    """Classic fourth-order Runge-Kutta. Fixed-step only: no error estimate."""

    order = 4
    fsal = False
    has_error_estimate = False

    def step(self, g, t, x, dt, f_0, project):
        k_1 = g(x, t) if f_0 is None else f_0
        k_2 = g(add_scaled(x, (0.5 * dt, k_1)), t + 0.5 * dt)
        k_3 = g(add_scaled(x, (0.5 * dt, k_2)), t + 0.5 * dt)
        k_4 = g(add_scaled(x, (dt, k_3)), t + dt)
        x_1 = project(
            add_scaled(
                x,
                (dt / 6.0, k_1),
                (dt / 3.0, k_2),
                (dt / 3.0, k_3),
                (dt / 6.0, k_4),
            )
        )
        return x_1, None, None

    def step_fixed(self, g, t, x, dt, f_0, project):
        return self.step(g, t, x, dt, f_0, project)

tinydiffeq.Tsit5 dataclass

Tsitouras 5(4) explicit Runge-Kutta with embedded error estimate.

FSAL: the last stage k_7 = g(x_1, t + dt) is the next step's first stage, so an accepted adaptive step costs six fresh evaluations. Note k_7 is evaluated at the projected accepted state, so the FSAL cache stays consistent with the state actually carried forward when project binds.

Source code in src/tinydiffeq/solvers.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class Tsit5:
    """Tsitouras 5(4) explicit Runge-Kutta with embedded error estimate.

    FSAL: the last stage k_7 = g(x_1, t + dt) is the next step's first stage,
    so an accepted adaptive step costs six fresh evaluations. Note k_7 is
    evaluated at the *projected* accepted state, so the FSAL cache stays
    consistent with the state actually carried forward when `project` binds.
    """

    order = 5
    fsal = True
    has_error_estimate = True

    def _step(self, g, t, x, dt, f_0, project, *, need_error):
        k_1 = g(x, t) if f_0 is None else f_0
        k_2 = g(add_scaled(x, (dt * A_21, k_1)), t + C_2 * dt)
        k_3 = g(
            add_scaled(x, (dt, weighted_sum((k_1, k_2), (A_31, A_32)))),
            t + C_3 * dt,
        )
        k_4 = g(
            add_scaled(x, (dt, weighted_sum((k_1, k_2, k_3), (A_41, A_42, A_43)))),
            t + C_4 * dt,
        )
        k_5 = g(
            add_scaled(
                x,
                (dt, weighted_sum((k_1, k_2, k_3, k_4), (A_51, A_52, A_53, A_54))),
            ),
            t + C_5 * dt,
        )
        k_6 = g(
            add_scaled(
                x,
                (
                    dt,
                    weighted_sum(
                        (k_1, k_2, k_3, k_4, k_5),
                        (A_61, A_62, A_63, A_64, A_65),
                    ),
                ),
            ),
            t + C_6 * dt,
        )
        x_1 = project(
            add_scaled(
                x,
                (
                    dt,
                    weighted_sum(
                        (k_1, k_2, k_3, k_4, k_5, k_6),
                        (B_1, B_2, B_3, B_4, B_5, B_6),
                    ),
                ),
            )
        )
        k_7 = g(x_1, t + C_7 * dt)
        if need_error:
            err = jax.tree.map(
                lambda value: dt * value,
                weighted_sum(
                    (k_1, k_2, k_3, k_4, k_5, k_6, k_7),
                    (E_1, E_2, E_3, E_4, E_5, E_6, E_7),
                ),
            )
        else:
            err = None
        return x_1, k_7, err

    def step(self, g, t, x, dt, f_0, project):
        return self._step(g, t, x, dt, f_0, project, need_error=True)

    def step_fixed(self, g, t, x, dt, f_0, project):
        """Take a fixed step without constructing the unused embedded error."""
        return self._step(g, t, x, dt, f_0, project, need_error=False)

step_fixed(g, t, x, dt, f_0, project)

Take a fixed step without constructing the unused embedded error.

Source code in src/tinydiffeq/solvers.py

def step_fixed(self, g, t, x, dt, f_0, project):
    """Take a fixed step without constructing the unused embedded error."""
    return self._step(g, t, x, dt, f_0, project, need_error=False)

tinydiffeq.Rodas5P dataclass

Fifth-order Rodas5P Rosenbrock--Wanner method.

An eight-stage, linearly implicit method with an embedded error estimate and a stiff-aware fourth-order continuous extension, supported by :func:tinydiffeq.solve_ode and :func:tinydiffeq.solve_semi_explicit_dae with one dense LU factorization reused across the stages of each attempted step. The implementation follows Steinebach (2023) and SciML's OrdinaryDiffEqRosenbrock.Rodas5P:

• https://doi.org/10.1007/s10543-023-00967-x
• https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock

Source code in src/tinydiffeq/solvers.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class Rodas5P:
    """Fifth-order Rodas5P Rosenbrock--Wanner method.

    An eight-stage, linearly implicit method with an embedded error estimate
    and a stiff-aware fourth-order continuous extension, supported by
    :func:`tinydiffeq.solve_ode` and
    :func:`tinydiffeq.solve_semi_explicit_dae` with one dense LU
    factorization reused across the stages of each attempted step. The
    implementation follows Steinebach (2023) and SciML's
    ``OrdinaryDiffEqRosenbrock.Rodas5P``:

    - https://doi.org/10.1007/s10543-023-00967-x
    - https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock
    """

    order = 5
    fsal = False
    has_error_estimate = True

tinydiffeq.EulerMaruyama dataclass

Euler-Maruyama for Ito SDEs with diagonal noise. Fixed-step only.

Strong order 0.5 for multiplicative noise. sample_noise returns the Brownian increments with the same pytree structure as the state and a leading n_steps axis.

Source code in src/tinydiffeq/solvers.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class EulerMaruyama:
    """Euler-Maruyama for Ito SDEs with diagonal noise. Fixed-step only.

    Strong order 0.5 for multiplicative noise. ``sample_noise`` returns the
    Brownian increments with the same pytree structure as the state and a
    leading ``n_steps`` axis.
    """

    order = 1
    strong_order = 0.5

    def sample_noise(self, x_0, key, n_steps, dt, dtype):
        return diagonal_brownian_increments(x_0, key, n_steps, dt, dtype)

    def step(self, g_drift, g_diffusion, t, x, dt, noise, project):
        return project(
            add_scaled(
                x, (dt, g_drift(x, t)), (1.0, multiply(g_diffusion(x, t), noise))
            )
        )

tinydiffeq.Milstein dataclass

Milstein for Ito SDEs with diagonal noise. Fixed-step only.

Strong order 1.0 under the diagonal commutativity condition: each diffusion component may depend only on its own state component. The correction (1/2) g g' (d_w^2 - dt) evaluates g g' as the forward-mode derivative of the diffusion field in the direction of its own value, which equals the diagonal term exactly in that case. sample_noise matches EulerMaruyama.

Source code in src/tinydiffeq/solvers.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class Milstein:
    """Milstein for Ito SDEs with diagonal noise. Fixed-step only.

    Strong order 1.0 under the diagonal commutativity condition: each
    diffusion component may depend only on its own state component. The
    correction ``(1/2) g g' (d_w^2 - dt)`` evaluates ``g g'`` as the
    forward-mode derivative of the diffusion field in the direction of its
    own value, which equals the diagonal term exactly in that case.
    ``sample_noise`` matches ``EulerMaruyama``.
    """

    order = 1
    strong_order = 1.0

    def sample_noise(self, x_0, key, n_steps, dt, dtype):
        return diagonal_brownian_increments(x_0, key, n_steps, dt, dtype)

    def step(self, g_drift, g_diffusion, t, x, dt, noise, project):
        g_value = g_diffusion(x, t)
        _, dg_g = jax.jvp(lambda state: g_diffusion(state, t), (x,), (g_value,))
        correction = jax.tree.map(lambda dg, w: 0.5 * dg * (w * w - dt), dg_g, noise)
        return project(
            add_scaled(
                x,
                (dt, g_drift(x, t)),
                (1.0, multiply(g_value, noise)),
                (1.0, correction),
            )
        )

tinydiffeq.SRA1 dataclass

Rossler SRA1 stochastic Runge-Kutta for Ito SDEs with additive diagonal noise. Fixed-step only.

Strong order 1.5 when the diffusion is independent of the state (it may depend on time). sample_noise returns (d_w, d_z): two independent sqrt(dt) * N(0, 1) draws per step. The time-Wiener integral I_10 / dt is realized internally as (d_w + d_z / sqrt(3)) / 2, reproducing its variance dt^3 / 3 and covariance dt^2 / 2 with the increment.

Source code in src/tinydiffeq/solvers.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class SRA1:
    """Rossler SRA1 stochastic Runge-Kutta for Ito SDEs with additive
    diagonal noise. Fixed-step only.

    Strong order 1.5 when the diffusion is independent of the state (it may
    depend on time). ``sample_noise`` returns ``(d_w, d_z)``: two independent
    ``sqrt(dt) * N(0, 1)`` draws per step. The time-Wiener integral
    ``I_10 / dt`` is realized internally as ``(d_w + d_z / sqrt(3)) / 2``,
    reproducing its variance ``dt^3 / 3`` and covariance ``dt^2 / 2`` with
    the increment.
    """

    order = 2
    strong_order = 1.5

    def sample_noise(self, x_0, key, n_steps, dt, dtype):
        key_w, key_z = jax.random.split(key)
        return (
            diagonal_brownian_increments(x_0, key_w, n_steps, dt, dtype),
            diagonal_brownian_increments(x_0, key_z, n_steps, dt, dtype),
        )

    def step(self, g_drift, g_diffusion, t, x, dt, noise, project):
        d_w, d_z = noise
        g_0 = g_diffusion(x, t)
        g_1 = g_diffusion(x, t + dt)
        chi = jax.tree.map(lambda w, z: 0.5 * (w + z * INV_SQRT_3), d_w, d_z)
        k_1 = g_drift(x, t)
        stage = add_scaled(x, (0.75 * dt, k_1), (1.5, multiply(g_1, chi)))
        k_2 = g_drift(stage, t + 0.75 * dt)
        return project(
            add_scaled(
                x,
                (dt / 3.0, k_1),
                (2.0 * dt / 3.0, k_2),
                (1.0, multiply(g_1, d_w)),
                (1.0, multiply(add_scaled(g_0, (-1.0, g_1)), chi)),
            )
        )

tinydiffeq.DenseExponential dataclass

Dense scaling-and-squaring exponential for a fixed linear operator.

Source code in src/tinydiffeq/exponential.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class DenseExponential:
    """Dense scaling-and-squaring exponential for a fixed linear operator."""

tinydiffeq.KrylovExponential dataclass

Static Arnoldi exponential action for an array or callable operator.

krylov_dim, num_substeps, and reorthogonalization_passes are static compilation controls. Two-pass classical Gram--Schmidt is the stable default; one pass reduces the dominant basis memory traffic when the operator has been validated for it. The precision-dependent default error tolerances are 1e-5/1e-7 for float32 and 1e-10/1e-12 for float64.

Source code in src/tinydiffeq/exponential.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class KrylovExponential:
    """Static Arnoldi exponential action for an array or callable operator.

    ``krylov_dim``, ``num_substeps``, and ``reorthogonalization_passes`` are
    static compilation controls. Two-pass classical Gram--Schmidt is the
    stable default; one pass reduces the dominant basis memory traffic when
    the operator has been validated for it. The precision-dependent default
    error tolerances are ``1e-5``/``1e-7`` for float32 and
    ``1e-10``/``1e-12`` for float64.
    """

    krylov_dim: int = field(default=30, metadata=dict(static=True))
    num_substeps: int = field(default=1, metadata=dict(static=True))
    reorthogonalization_passes: int = field(default=2, metadata=dict(static=True))
    rtol: float | None = None
    atol: float | None = None

    def __post_init__(self):
        if not isinstance(self.krylov_dim, int) or isinstance(self.krylov_dim, bool):
            raise TypeError("KrylovExponential.krylov_dim must be a positive int")
        if self.krylov_dim < 1:
            raise ValueError("KrylovExponential.krylov_dim must be a positive int")
        if not isinstance(self.num_substeps, int) or isinstance(
            self.num_substeps, bool
        ):
            raise TypeError("KrylovExponential.num_substeps must be a positive int")
        if self.num_substeps < 1:
            raise ValueError("KrylovExponential.num_substeps must be a positive int")
        if (
            not isinstance(self.reorthogonalization_passes, int)
            or isinstance(self.reorthogonalization_passes, bool)
            or self.reorthogonalization_passes not in (1, 2)
        ):
            raise ValueError(
                "KrylovExponential.reorthogonalization_passes must be 1 or 2"
            )
        for name, value in (("rtol", self.rtol), ("atol", self.atol)):
            if value is not None and value < 0:
                raise ValueError(f"KrylovExponential.{name} must be nonnegative")

tinydiffeq.AdaptiveKrylovExponential dataclass

Adaptive matrix-free Arnoldi exponential action.

The Krylov dimension remains static while accepted internal time slices adapt to the leading-term Arnoldi residual. max_steps bounds all accepted and rejected attempts, keeping compiled shapes static. The precision-dependent tolerance defaults match KrylovExponential.

Source code in src/tinydiffeq/exponential.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class AdaptiveKrylovExponential:
    """Adaptive matrix-free Arnoldi exponential action.

    The Krylov dimension remains static while accepted internal time slices
    adapt to the leading-term Arnoldi residual. ``max_steps`` bounds all
    accepted and rejected attempts, keeping compiled shapes static. The
    precision-dependent tolerance defaults match ``KrylovExponential``.
    """

    krylov_dim: int = field(default=30, metadata=dict(static=True))
    max_steps: int = field(default=128, metadata=dict(static=True))
    reorthogonalization_passes: int = field(default=2, metadata=dict(static=True))
    initial_step: float | None = None
    safety: float = 0.9
    min_factor: float = 0.2
    max_factor: float = 5.0
    rtol: float | None = None
    atol: float | None = None

    def __post_init__(self):
        for name, value in (
            ("krylov_dim", self.krylov_dim),
            ("max_steps", self.max_steps),
        ):
            if not isinstance(value, int) or isinstance(value, bool):
                raise TypeError(
                    f"AdaptiveKrylovExponential.{name} must be a positive int"
                )
            if value < 1:
                raise ValueError(
                    f"AdaptiveKrylovExponential.{name} must be a positive int"
                )
        if (
            not isinstance(self.reorthogonalization_passes, int)
            or isinstance(self.reorthogonalization_passes, bool)
            or self.reorthogonalization_passes not in (1, 2)
        ):
            raise ValueError(
                "AdaptiveKrylovExponential.reorthogonalization_passes must be 1 or 2"
            )
        if self.initial_step is not None and self.initial_step <= 0:
            raise ValueError("AdaptiveKrylovExponential.initial_step must be positive")
        if not 0 < self.safety <= 1:
            raise ValueError("AdaptiveKrylovExponential.safety must be in (0, 1]")
        if not 0 < self.min_factor < 1:
            raise ValueError("AdaptiveKrylovExponential.min_factor must be in (0, 1)")
        if self.max_factor <= 1:
            raise ValueError(
                "AdaptiveKrylovExponential.max_factor must be greater than 1"
            )
        for name, value in (("rtol", self.rtol), ("atol", self.atol)):
            if value is not None and value < 0:
                raise ValueError(
                    f"AdaptiveKrylovExponential.{name} must be nonnegative"
                )

Step-size controllers

tinydiffeq.ConstantStepSize dataclass

Accept every step and keep the carried step size unchanged.

Source code in src/tinydiffeq/controllers.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class ConstantStepSize:
    """Accept every step and keep the carried step size unchanged."""

    uses_error_estimate = False

    def init(self, x_0):
        return ()

    def adapt(self, x_0, x_1, err, dt_used, dt_prev, order, state, time_scale=1.0):
        return jnp.asarray(True), dt_prev, state

tinydiffeq.IController dataclass

Integral step-size controller with max-norm error.

Accept iff E = max(|err| / (atol + rtol * max(|x_0|, |x_1|))) <= 1 (forced accept once the step reaches dt_min), and propose dt_next = dt_used * clip(safety * E**(-1/order), factor_min, factor_max) clipped to [dt_min, dt_max]. Omitted rtol/atol default to 1e-4/1e-6 for float32 states and 1e-7/1e-9 for float64; dt_min defaults to ten machine epsilons in the time dtype, scaled by max(1, |t_1|).

Source code in src/tinydiffeq/controllers.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class IController:
    """Integral step-size controller with max-norm error.

    Accept iff ``E = max(|err| / (atol + rtol * max(|x_0|, |x_1|))) <= 1``
    (forced accept once the step reaches ``dt_min``), and propose
    ``dt_next = dt_used * clip(safety * E**(-1/order), factor_min, factor_max)``
    clipped to ``[dt_min, dt_max]``. Omitted ``rtol``/``atol`` default to
    ``1e-4``/``1e-6`` for float32 states and ``1e-7``/``1e-9`` for float64;
    ``dt_min`` defaults to ten machine epsilons in the time dtype, scaled by
    ``max(1, |t_1|)``.
    """

    rtol: float | None = None
    atol: float | None = None
    dt_min: float | None = None
    dt_max: float = float("inf")
    safety: float = SAFETY
    factor_min: float = MIN_FACTOR
    factor_max: float = MAX_FACTOR

    uses_error_estimate = True

    def init(self, x_0):
        return ()

    def adapt(self, x_0, x_1, err, dt_used, dt_prev, order, state, time_scale=1.0):
        dtype = jax.tree.leaves(x_0)[0].dtype
        rtol, atol = _resolve_tolerances(self.rtol, self.atol, dtype)
        dt_min = _resolve_dt_min(
            self.dt_min, jnp.result_type(dt_used), jnp.asarray(time_scale)
        )
        dt_max = jnp.asarray(self.dt_max, jnp.result_type(dt_used))
        safety = jnp.asarray(self.safety, dtype)
        factor_min = jnp.asarray(self.factor_min, dtype)
        factor_max = jnp.asarray(self.factor_max, dtype)
        # The controller is wrapped in stop_gradient: accept/reject is a
        # non-differentiable branch either way, the gradient of E**(-1/order)
        # blows up at the exact-zero error of a flat-start policy, and the
        # d(dt)/dtheta term only slides sample points along the visited
        # trajectory -- irrelevant to a residual that must vanish at every
        # state. The states themselves remain fully differentiable through
        # the solver stages.
        scaled_error = jnp.asarray(
            jax.lax.stop_gradient(error_ratio(x_0, x_1, err, rtol, atol)),
            dtype,
        )
        accept = (scaled_error <= 1.0) | (dt_used <= dt_min)
        error_floor = jnp.asarray(jnp.finfo(dtype).eps, dtype)
        factor = jnp.clip(
            safety * jnp.maximum(scaled_error, error_floor) ** (-1.0 / order),
            factor_min,
            factor_max,
        )
        dt_next = jnp.clip(jax.lax.stop_gradient(dt_used) * factor, dt_min, dt_max)
        return accept, dt_next, state

tinydiffeq.PIController dataclass

Proportional-integral step-size controller with max-norm error.

In addition to the current scaled error E, this controller carries the previous accepted step's error E_prev (starting at one) and proposes dt_next = dt_used * clip(safety * E**(-(p_coeff+i_coeff)/order) * E_prev**(p_coeff/order), factor_min, factor_max). The defaults p_coeff=0.4 and i_coeff=0.3 damp step-size oscillations; p_coeff=0, i_coeff=1 reproduces :class:IController, and the tolerance and dt_min defaults match it.

Source code in src/tinydiffeq/controllers.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class PIController:
    """Proportional-integral step-size controller with max-norm error.

    In addition to the current scaled error ``E``, this controller carries the
    previous accepted step's error ``E_prev`` (starting at one) and proposes
    ``dt_next = dt_used * clip(safety * E**(-(p_coeff+i_coeff)/order)``
    ``* E_prev**(p_coeff/order), factor_min, factor_max)``. The defaults
    ``p_coeff=0.4`` and ``i_coeff=0.3`` damp step-size oscillations;
    ``p_coeff=0, i_coeff=1`` reproduces :class:`IController`, and the
    tolerance and ``dt_min`` defaults match it.
    """

    rtol: float | None = None
    atol: float | None = None
    p_coeff: float = 0.4
    i_coeff: float = 0.3
    dt_min: float | None = None
    dt_max: float = float("inf")
    safety: float = SAFETY
    factor_min: float = MIN_FACTOR
    factor_max: float = MAX_FACTOR

    uses_error_estimate = True

    def init(self, x_0):
        return jnp.asarray(1.0, jax.tree.leaves(x_0)[0].dtype)

    def adapt(self, x_0, x_1, err, dt_used, dt_prev, order, state, time_scale=1.0):
        dtype = jax.tree.leaves(x_0)[0].dtype
        rtol, atol = _resolve_tolerances(self.rtol, self.atol, dtype)
        dt_min = _resolve_dt_min(
            self.dt_min, jnp.result_type(dt_used), jnp.asarray(time_scale)
        )
        dt_max = jnp.asarray(self.dt_max, jnp.result_type(dt_used))
        safety = jnp.asarray(self.safety, dtype)
        factor_min = jnp.asarray(self.factor_min, dtype)
        factor_max = jnp.asarray(self.factor_max, dtype)
        p_coeff = jnp.asarray(self.p_coeff, dtype)
        i_coeff = jnp.asarray(self.i_coeff, dtype)
        scaled_error = jnp.asarray(
            jax.lax.stop_gradient(error_ratio(x_0, x_1, err, rtol, atol)),
            dtype,
        )
        accept = (scaled_error <= 1.0) | (dt_used <= dt_min)
        error_floor = jnp.asarray(jnp.finfo(dtype).eps, dtype)
        safe_error_ratio = jnp.maximum(scaled_error, error_floor)
        safe_previous_error_ratio = jnp.maximum(state, error_floor)
        factor = jnp.clip(
            safety
            * safe_error_ratio ** (-(p_coeff + i_coeff) / order)
            * safe_previous_error_ratio ** (p_coeff / order),
            factor_min,
            factor_max,
        )
        dt_next = jnp.clip(jax.lax.stop_gradient(dt_used) * factor, dt_min, dt_max)
        state_next = jnp.where(accept, scaled_error, state)
        return accept, dt_next, state_next

Markov-chain models and methods

tinydiffeq.DiscreteMarkovChain

Prepared dense homogeneous finite-state transition matrix.

Construction validates and normalizes the rows and builds Vose alias tables. Construct once outside transformed code, then pass the object through jit or as a shared vmap argument.

Source code in src/tinydiffeq/markov.py

@jax.tree_util.register_pytree_node_class
class DiscreteMarkovChain:
    """Prepared dense homogeneous finite-state transition matrix.

    Construction validates and normalizes the rows and builds Vose alias tables.
    Construct once outside transformed code, then pass the object through ``jit``
    or as a shared ``vmap`` argument.
    """

    def __init__(self, transition_matrix):
        value, host = _concrete_square_matrix(transition_matrix, "transition_matrix")
        normalized = _normalize_transition_rows(host, "transition_matrix")
        self.transition_matrix = jnp.asarray(normalized, dtype=value.dtype)
        self.alias_probability, self.alias_index = _alias_tables(
            normalized, value.dtype
        )

    @property
    def num_states(self):
        return self.transition_matrix.shape[0]

    def tree_flatten(self):
        return (
            self.transition_matrix,
            self.alias_probability,
            self.alias_index,
        ), None

    @classmethod
    def tree_unflatten(cls, auxiliary, children):
        instance = object.__new__(cls)
        (
            instance.transition_matrix,
            instance.alias_probability,
            instance.alias_index,
        ) = children
        return instance

tinydiffeq.ContinuousTimeMarkovChain

Prepared dense homogeneous finite-state generator matrix.

Off-diagonal entries must be nonnegative and each row must sum to zero. Zero rows are absorbing states. Construction extracts exit rates and builds alias tables for the embedded jump chain; construct outside transformed code.

Source code in src/tinydiffeq/markov.py

@jax.tree_util.register_pytree_node_class
class ContinuousTimeMarkovChain:
    """Prepared dense homogeneous finite-state generator matrix.

    Off-diagonal entries must be nonnegative and each row must sum to zero.
    Zero rows are absorbing states. Construction extracts exit rates and builds
    alias tables for the embedded jump chain; construct outside transformed code.
    """

    def __init__(self, generator):
        value, host = _concrete_square_matrix(generator, "generator")
        num_states = value.shape[0]
        tolerance = 100.0 * float(jnp.finfo(value.dtype).eps)
        embedded = []
        rates = []
        for row_index, row in enumerate(host):
            if any(not math.isfinite(entry) for entry in row):
                raise ValueError("generator must contain only finite values")
            off_diagonal = [
                entry if column != row_index else 0.0
                for column, entry in enumerate(row)
            ]
            if any(entry < 0.0 for entry in off_diagonal):
                raise ValueError("generator off-diagonal entries must be nonnegative")
            rate = sum(off_diagonal)
            scale = max(1.0, rate, abs(row[row_index]))
            if abs(row[row_index] + rate) > tolerance * scale:
                raise ValueError("generator rows must sum to zero")
            rates.append(rate)
            if rate == 0.0:
                probabilities = [0.0] * num_states
                probabilities[row_index] = 1.0
            else:
                probabilities = [entry / rate for entry in off_diagonal]
            embedded.append(probabilities)
        self.generator = value
        self.exit_rates = jnp.asarray(rates, dtype=value.dtype)
        self.alias_probability, self.alias_index = _alias_tables(embedded, value.dtype)

    @property
    def num_states(self):
        return self.generator.shape[0]

    def tree_flatten(self):
        return (
            self.generator,
            self.exit_rates,
            self.alias_probability,
            self.alias_index,
        ), None

    @classmethod
    def tree_unflatten(cls, auxiliary, children):
        instance = object.__new__(cls)
        (
            instance.generator,
            instance.exit_rates,
            instance.alias_probability,
            instance.alias_index,
        ) = children
        return instance

tinydiffeq.MatrixFreeContinuousTimeMarkovChain

Fixed CTMC forward generator represented by a pytree linear action.

forward_generator(probabilities) must return the same probability-pytree structure and dtype and represent the forward equation dπ/dt = L(π). The callable is static JAX structure; close only over fixed model data or put changing arrays inside a callable pytree.

Source code in src/tinydiffeq/markov.py

@jax.tree_util.register_pytree_node_class
class MatrixFreeContinuousTimeMarkovChain:
    """Fixed CTMC forward generator represented by a pytree linear action.

    ``forward_generator(probabilities)`` must return the same probability-pytree
    structure and dtype and represent the forward equation ``dπ/dt = L(π)``.
    The callable is static JAX structure; close only over fixed model data or put
    changing arrays inside a callable pytree.
    """

    def __init__(self, forward_generator):
        if not callable(forward_generator):
            raise TypeError("forward_generator must be callable")
        self.forward_generator = forward_generator

    def tree_flatten(self):
        return (), self.forward_generator

    @classmethod
    def tree_unflatten(cls, forward_generator, children):
        instance = object.__new__(cls)
        instance.forward_generator = forward_generator
        return instance

tinydiffeq.SequentialMarkov dataclass

Chronological scan method; unroll=1 is the CPU-oriented default.

Source code in src/tinydiffeq/markov.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class SequentialMarkov:
    """Chronological scan method; ``unroll=1`` is the CPU-oriented default."""

    unroll: int = field(default=1, metadata=dict(static=True))

    def __post_init__(self):
        if not isinstance(self.unroll, int) or isinstance(self.unroll, bool):
            raise TypeError("SequentialMarkov.unroll must be a positive int")
        if self.unroll < 1:
            raise ValueError("SequentialMarkov.unroll must be a positive int")

tinydiffeq.AssociativeMarkov dataclass

Parallel-prefix random-map composition method.

Source code in src/tinydiffeq/markov.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class AssociativeMarkov:
    """Parallel-prefix random-map composition method."""

tinydiffeq.MatrixPowerMarkov dataclass

Binary matrix powering for a DTMC distribution at one endpoint.

Source code in src/tinydiffeq/markov.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class MatrixPowerMarkov:
    """Binary matrix powering for a DTMC distribution at one endpoint."""

Output selection and results

tinydiffeq.SaveAt dataclass

What the solve functions return. Exactly one mode must be set.

t_1=True (the solver default) returns the endpoint only. ts=grid interpolates the internal steps onto a fixed query grid — output shape is (len(ts), ...) however many steps the controller takes, and ts is a data leaf; with exact=True an explicit constant-step ODE instead gathers states at queries that must coincide with realized knots. steps=True returns the initial state and accepted steps as the valid prefix of a max_steps + 1 buffer, padded with the last valid row (fill="last") or inf (fill="inf") and masked by Solution.accepted.

Source code in src/tinydiffeq/save_at.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class SaveAt:
    """What the solve functions return. Exactly one mode must be set.

    ``t_1=True`` (the solver default) returns the endpoint only. ``ts=grid``
    interpolates the internal steps onto a fixed query grid — output shape is
    ``(len(ts), ...)`` however many steps the controller takes, and ``ts`` is
    a data leaf; with ``exact=True`` an explicit constant-step ODE instead
    gathers states at queries that must coincide with realized knots.
    ``steps=True`` returns the initial state and accepted steps as the valid
    prefix of a ``max_steps + 1`` buffer, padded with the last valid row
    (``fill="last"``) or ``inf`` (``fill="inf"``) and masked by
    ``Solution.accepted``.
    """

    t_1: bool = field(default=False, metadata=dict(static=True))
    ts: ArrayLike | None = None
    steps: bool = field(default=False, metadata=dict(static=True))
    fill: str = field(default="last", metadata=dict(static=True))
    exact: bool = field(default=False, metadata=dict(static=True))

    def __post_init__(self):
        modes = int(bool(self.t_1)) + int(self.ts is not None) + int(bool(self.steps))
        if modes != 1:
            raise ValueError(
                "SaveAt requires exactly one of t_1=True, ts=..., steps=True"
            )
        if self.fill not in ("last", "inf"):
            raise ValueError('SaveAt fill must be "last" or "inf"')
        if self.exact and self.ts is None:
            raise ValueError("SaveAt exact=True requires ts=...")

tinydiffeq.LMRootSolver dataclass

Configuration for algebraic root solves in semi-explicit DAEs/SDAEs.

The implementation is :class:nlls_gram.LevenbergMarquardt at its defaults. max_steps bounds one algebraic root's nonlinear iterations, independently of the integration's time-step budget. Roots use residual stopping only: gtol and xtol must remain zero, and every accepted root must report CONVERGED with Euclidean residual norm strictly below atol (None selects 1e-6 in float32, 1e-10 in float64). solver_options is a mapping (or pairs) forwarded verbatim to the LevenbergMarquardt constructor, e.g. solver_options={"linear_solver": QR()}; cache_jacobian and geodesic_acceleration are fixed to False. predictor selects the explicit-stage warm start: "previous" reuses the most recent successful root, "secant" extrapolates it to a strictly later stage time and assumes the continued root is locally unique.

Source code in src/tinydiffeq/dae.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class LMRootSolver:
    """Configuration for algebraic root solves in semi-explicit DAEs/SDAEs.

    The implementation is :class:`nlls_gram.LevenbergMarquardt` at its
    defaults. ``max_steps`` bounds one algebraic root's nonlinear iterations,
    independently of the integration's time-step budget. Roots use residual
    stopping only: ``gtol`` and ``xtol`` must remain zero, and every accepted
    root must report ``CONVERGED`` with Euclidean residual norm strictly below
    ``atol`` (``None`` selects ``1e-6`` in float32, ``1e-10`` in float64).
    ``solver_options`` is a mapping (or pairs) forwarded verbatim to the
    ``LevenbergMarquardt`` constructor, e.g.
    ``solver_options={"linear_solver": QR()}``; ``cache_jacobian`` and
    ``geodesic_acceleration`` are fixed to ``False``. ``predictor`` selects
    the explicit-stage warm start: ``"previous"`` reuses the most recent
    successful root, ``"secant"`` extrapolates it to a strictly later stage
    time and assumes the continued root is locally unique.
    """

    max_steps: int = field(default=8, metadata=dict(static=True))
    atol: float | None = field(default=None, metadata=dict(static=True))
    gtol: float = field(default=0.0, metadata=dict(static=True))
    xtol: float = field(default=0.0, metadata=dict(static=True))
    solver_options: Any = field(default=(), metadata=dict(static=True))
    predictor: str = field(default="previous", metadata=dict(static=True))

    def __post_init__(self):
        if not isinstance(self.max_steps, int) or isinstance(self.max_steps, bool):
            raise ValueError("LMRootSolver.max_steps must be a positive int")
        if self.max_steps <= 0:
            raise ValueError("LMRootSolver.max_steps must be a positive int")
        if self.atol is not None and self.atol <= 0:
            raise ValueError("LMRootSolver.atol must be positive or None")
        if self.gtol != 0:
            raise ValueError("LMRootSolver.gtol must be zero for DAE root solves")
        if self.xtol != 0:
            raise ValueError("LMRootSolver.xtol must be zero for DAE root solves")
        if not isinstance(self.predictor, str):
            raise TypeError("LMRootSolver.predictor must be a string")
        if self.predictor not in ("previous", "secant"):
            raise ValueError(
                'LMRootSolver.predictor must be either "previous" or "secant"'
            )
        try:
            options = tuple(sorted(dict(self.solver_options).items()))
        except (TypeError, ValueError) as error:
            raise TypeError(
                "LMRootSolver.solver_options must be a mapping or key/value pairs"
            ) from error
        fixed = _FIXED_SOLVER_OPTIONS.intersection(name for name, _ in options)
        if fixed:
            raise ValueError(
                "LMRootSolver fixes " + ", ".join(sorted(fixed)) + " for algebraic "
                "roots; they cannot be set through solver_options"
            )
        object.__setattr__(self, "solver_options", options)

tinydiffeq.DAESolution dataclass

Result of the deterministic or stochastic semi-explicit DAE solvers.

ts/ys/zs follow the :class:Solution shape contract, with zs the algebraic states. Explicit-method saved values sit at converged roots; Rodas5P satisfies the constraint to integration accuracy after its initial consistency root, and requested-grid interpolants need not satisfy it exactly. num_root_solves counts logical active root calls (including failures and the initial consistency solve) and num_root_steps sums their LM update steps; like num_steps, both are path diagnostics with exact-zero tangents.

Source code in src/tinydiffeq/solution.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class DAESolution:
    """Result of the deterministic or stochastic semi-explicit DAE solvers.

    ``ts``/``ys``/``zs`` follow the :class:`Solution` shape contract, with
    ``zs`` the algebraic states. Explicit-method saved values sit at
    converged roots; Rodas5P satisfies the constraint to integration accuracy
    after its initial consistency root, and requested-grid interpolants need
    not satisfy it exactly. ``num_root_solves`` counts logical active root
    calls (including failures and the initial consistency solve) and
    ``num_root_steps`` sums their LM update steps; like ``num_steps``, both
    are path diagnostics with exact-zero tangents.
    """

    ts: jax.Array
    ys: Any
    zs: Any
    ok: jax.Array
    num_accepted: jax.Array
    accepted: jax.Array | None = None
    aux: Any = None
    num_steps: jax.Array | None = None
    num_root_solves: jax.Array | None = None
    num_root_steps: jax.Array | None = None

tinydiffeq.BVPSolution dataclass

Result of solve_bvp.

Arrays are padded to the static max_nodes: the t tail repeats the right endpoint and the y/yp tails repeat the last active row, so sol(ts) evaluates exactly the C1 cubic spline scipy's solve_bvp returns and sol.derivative(ts) its derivative. z holds the solved unknown parameters (None when the problem has none), rms_residuals is zero on inactive intervals, num_nodes counts active mesh nodes, and num_iterations is scipy's niter. status uses scipy's codes (0 converged, 1 max_nodes exceeded, 2 singular Jacobian, 3 boundary-condition tolerance unsatisfied) and ok is status == 0; a failed status returns the last iterate, which may be non-finite. Under AD only y, yp, z, and aux carry tangents with respect to p; every other field is differentiation-inert with exact-zero tangents.

Source code in src/tinydiffeq/solution.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class BVPSolution:
    """Result of ``solve_bvp``.

    Arrays are padded to the static ``max_nodes``: the ``t`` tail repeats the
    right endpoint and the ``y``/``yp`` tails repeat the last active row, so
    ``sol(ts)`` evaluates exactly the C1 cubic spline scipy's ``solve_bvp``
    returns and ``sol.derivative(ts)`` its derivative. ``z`` holds the solved
    unknown parameters (``None`` when the
    problem has none), ``rms_residuals`` is zero on inactive intervals,
    ``num_nodes`` counts active mesh nodes, and ``num_iterations`` is scipy's
    ``niter``. ``status`` uses scipy's codes (0 converged, 1 ``max_nodes``
    exceeded, 2 singular Jacobian, 3 boundary-condition tolerance unsatisfied)
    and ``ok`` is ``status == 0``; a failed status returns the last iterate,
    which may be non-finite. Under AD only ``y``, ``yp``, ``z``, and ``aux``
    carry tangents with respect to ``p``; every other field is
    differentiation-inert with exact-zero tangents.
    """

    t: jax.Array
    y: Any
    yp: Any
    z: Any
    rms_residuals: jax.Array
    num_nodes: jax.Array
    num_iterations: jax.Array
    status: jax.Array
    ok: jax.Array
    aux: Any = None

    def __call__(self, ts):
        return hermite_interpolate(ts, self.t, self.y, self.yp)

    def derivative(self, ts):
        return hermite_derivative(ts, self.t, self.y, self.yp)

tinydiffeq.Solution dataclass

Result of solve_ode/solve_sde/solve_linear_ode.

ts/xs hold times and states in the shape dictated by SaveAt. ok is a scalar bool: the integration reached t_1 and every required saved output was valid. Outputs are never poisoned; callers that want diverging values map jnp.where(sol.ok, x, jnp.inf) over leaves. num_accepted counts accepted steps, num_steps counts logical attempts including rejections, accepted masks the valid prefix in steps mode (row 0 is always True), and aux holds the field's saved auxiliary pytree with the same leading saved-time axis as xs.

Source code in src/tinydiffeq/solution.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class Solution:
    """Result of ``solve_ode``/``solve_sde``/``solve_linear_ode``.

    ``ts``/``xs`` hold times and states in the shape dictated by ``SaveAt``.
    ``ok`` is a scalar bool: the integration reached ``t_1`` and every
    required saved output was valid. Outputs are never poisoned; callers that
    want diverging values map ``jnp.where(sol.ok, x, jnp.inf)`` over leaves.
    ``num_accepted`` counts accepted steps, ``num_steps`` counts logical
    attempts including rejections, ``accepted`` masks the valid prefix in
    ``steps`` mode (row 0 is always True), and ``aux`` holds the field's
    saved auxiliary pytree with the same leading saved-time axis as ``xs``.
    """

    ts: jax.Array
    xs: Any
    ok: jax.Array
    num_accepted: jax.Array
    accepted: jax.Array | None = None
    aux: Any = None
    num_steps: jax.Array | None = None

tinydiffeq.MarkovDistribution dataclass

Forecast probability mass, evaluation steps/times, and validity flag.

Source code in src/tinydiffeq/markov.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class MarkovDistribution:
    """Forecast probability mass, evaluation steps/times, and validity flag."""

    ts: jax.Array
    probabilities: Any
    ok: jax.Array

Utilities

tinydiffeq.diagonal_brownian_increments(x_0, key, n_steps, dt, dtype)

Draw n_steps diagonal Brownian increments sqrt(dt) * N(0, 1).

Arrays retain the exact (n_steps,) + x_0.shape draw. Pytree states use one shared flat draw, partitioned into leaves in JAX's deterministic pytree leaf order.

Source code in src/tinydiffeq/solvers.py

def diagonal_brownian_increments(x_0, key, n_steps, dt, dtype):
    """Draw ``n_steps`` diagonal Brownian increments ``sqrt(dt) * N(0, 1)``.

    Arrays retain the exact ``(n_steps,) + x_0.shape`` draw. Pytree states use
    one shared flat draw, partitioned into leaves in JAX's deterministic
    pytree leaf order.
    """
    leaves, treedef = jax.tree.flatten(x_0)
    if treedef == jax.tree.structure(0):
        return jnp.sqrt(dt) * jax.random.normal(
            key, (n_steps,) + x_0.shape, dtype=dtype
        )
    sizes = [leaf.size for leaf in leaves]
    flat_noise = jnp.sqrt(dt) * jax.random.normal(
        key, (n_steps, sum(sizes)), dtype=dtype
    )
    noise_leaves = []
    start = 0
    for leaf, size in zip(leaves, sizes, strict=True):
        noise_leaves.append(
            flat_noise[:, start : start + size].reshape((n_steps,) + leaf.shape)
        )
        start += size
    return jax.tree.unflatten(treedef, noise_leaves)

tinydiffeq.hermite_interpolate(ts_query, knot_ts, knot_xs, knot_fs)

Cubic Hermite interpolation of (knot_ts, knot_xs, knot_fs) at ts_query, where knot_fs holds the time derivatives at the knots.

knot_ts must be nondecreasing; duplicate knots are allowed and queries falling on a zero-width bracket return the left knot value. Queries outside the knot span clamp to the boundary knot values (flat extrapolation) rather than evaluating the cubic outside its bracket.

Source code in src/tinydiffeq/interpolation.py

def hermite_interpolate(ts_query, knot_ts, knot_xs, knot_fs):
    """Cubic Hermite interpolation of ``(knot_ts, knot_xs, knot_fs)`` at
    ``ts_query``, where ``knot_fs`` holds the time derivatives at the knots.

    ``knot_ts`` must be nondecreasing; duplicate knots are allowed and
    queries falling on a zero-width bracket return the left knot value.
    Queries outside the knot span clamp to the boundary knot values (flat
    extrapolation) rather than evaluating the cubic outside its bracket.
    """
    n = knot_ts.shape[0]
    idx = jnp.clip(jnp.searchsorted(knot_ts, ts_query, side="right") - 1, 0, n - 2)
    t_left, t_right = knot_ts[idx], knot_ts[idx + 1]
    width = t_right - t_left
    degenerate = width <= 0.0
    # double-where: divide by the safe width BEFORE branching so neither the
    # primal nor its jvp/vjp ever sees a 0/0.
    width_safe = jnp.where(degenerate, 1.0, width)
    s = jnp.clip((ts_query - t_left) / width_safe, 0.0, 1.0)

    def interpolate_leaf(xs, fs):
        x_left, x_right = xs[idx], xs[idx + 1]
        f_left, f_right = fs[idx], fs[idx + 1]
        extra = xs.ndim - 1

        def bc(a):
            return a.reshape(a.shape + (1,) * extra)

        # Each leaf keeps its own dtype. This matters for mixed-precision aux
        # pytrees: a float64 time grid must not widen a float32 output leaf.
        s_leaf = s.astype(xs.dtype)
        width_leaf = width_safe.astype(xs.dtype)
        s_, w_, deg_ = bc(s_leaf), bc(width_leaf), bc(degenerate)
        h_00 = (1.0 + 2.0 * s_) * (1.0 - s_) ** 2
        h_10 = s_ * (1.0 - s_) ** 2
        h_01 = s_**2 * (3.0 - 2.0 * s_)
        h_11 = s_**2 * (s_ - 1.0)
        value = (
            h_00 * x_left + h_10 * w_ * f_left + h_01 * x_right + h_11 * w_ * f_right
        )
        return jnp.where(deg_, x_left, value)

    return jax.tree.map(interpolate_leaf, knot_xs, knot_fs)

tinydiffeq.hermite_derivative(ts_query, knot_ts, knot_xs, knot_fs)

Derivative of the C1 cubic Hermite interpolant at ts_query.

Uses searchsorted(side="left") so a query at a repeated right-endpoint knot lands on the last positive-width bracket and returns the knot derivative instead of the degenerate bracket's zero. Queries outside the knot span return zero, the derivative of the clamped value extension.

Source code in src/tinydiffeq/interpolation.py

def hermite_derivative(ts_query, knot_ts, knot_xs, knot_fs):
    """Derivative of the C1 cubic Hermite interpolant at ``ts_query``.

    Uses ``searchsorted(side="left")`` so a query at a repeated right-endpoint
    knot lands on the last positive-width bracket and returns the knot
    derivative instead of the degenerate bracket's zero. Queries outside the
    knot span return zero, the derivative of the clamped value extension.
    """
    n = knot_ts.shape[0]
    idx = jnp.clip(jnp.searchsorted(knot_ts, ts_query, side="left") - 1, 0, n - 2)
    t_left, t_right = knot_ts[idx], knot_ts[idx + 1]
    width = t_right - t_left
    degenerate = width <= 0.0
    width_safe = jnp.where(degenerate, 1.0, width)
    s = jnp.clip((ts_query - t_left) / width_safe, 0.0, 1.0)
    outside = (ts_query < knot_ts[0]) | (ts_query > knot_ts[n - 1])

    def derivative_leaf(xs, fs):
        x_left, x_right = xs[idx], xs[idx + 1]
        f_left, f_right = fs[idx], fs[idx + 1]
        extra = xs.ndim - 1

        def bc(a):
            return a.reshape(a.shape + (1,) * extra)

        s_leaf = s.astype(xs.dtype)
        width_leaf = width_safe.astype(xs.dtype)
        s_, w_, deg_, out_ = bc(s_leaf), bc(width_leaf), bc(degenerate), bc(outside)
        d00 = 6.0 * s_**2 - 6.0 * s_
        d10 = 3.0 * s_**2 - 4.0 * s_ + 1.0
        d11 = 3.0 * s_**2 - 2.0 * s_
        value = d00 * (x_left - x_right) / w_ + d10 * f_left + d11 * f_right
        value = jnp.where(deg_, f_left, value)
        return jnp.where(out_, jnp.zeros_like(value), value)

    return jax.tree.map(derivative_leaf, knot_xs, knot_fs)

tinydiffeq.cumulative_trapezoid(g, ts, *, substeps=1)

Cumulative composite-trapezoid integral of a time-only g(t) on the (possibly nonuniform) sorted grid ts.

Each grid interval is subdivided into substeps uniform panels, so the quadrature error shrinks with substeps without changing the output grid. Returns (integral, values) where integral[k] approximates the integral of g from ts[0] to ts[k] (integral[0] = 0) and values = g(ts); g may return any array shape, which is appended to the leading grid axis.

Source code in src/tinydiffeq/quadrature.py

def cumulative_trapezoid(g, ts, *, substeps=1):
    """Cumulative composite-trapezoid integral of a time-only ``g(t)`` on the
    (possibly nonuniform) sorted grid ``ts``.

    Each grid interval is subdivided into ``substeps`` uniform panels, so the
    quadrature error shrinks with ``substeps`` without changing the output
    grid. Returns ``(integral, values)`` where ``integral[k]`` approximates
    the integral of ``g`` from ``ts[0]`` to ``ts[k]`` (``integral[0] = 0``)
    and ``values = g(ts)``; ``g`` may return any array shape, which is
    appended to the leading grid axis.
    """
    if substeps < 1:
        raise ValueError("substeps must be at least 1")
    values = jax.vmap(g)(ts)

    def interval_increment(left, right):
        nodes = jnp.linspace(left, right, substeps + 1)
        node_values = jax.vmap(g)(nodes)
        widths = nodes[1:] - nodes[:-1]
        widths = widths.reshape(widths.shape + (1,) * (node_values.ndim - 1))
        return jnp.sum(0.5 * widths * (node_values[:-1] + node_values[1:]), axis=0)

    increments = jax.vmap(interval_increment)(ts[:-1], ts[1:])
    integral = jnp.concatenate(
        [jnp.zeros_like(increments[:1]), jnp.cumsum(increments, axis=0)]
    )
    return integral, values