Extend observer outputs and lazy context to linear and matrix-free solvers

[jlperla]

Motivation

The differential-equation solvers now expose one consistent saved-output contract: a differential field owns a real-floating aux pytree, while semi-explicit algebraic aux is internal context shared with the dynamics. Linear exponential solves and Markov distribution forecasts do not yet have an equivalent way to return derived moments without adding them to the propagated state.

Future matrix-free implicit ODE/DAE methods also need care: a Krylov residual/JVP operation should not materialize an expensive algebraic context unless the differential field or an observer actually consumes it.

Proposed scope

Universal observer output

Design a static, optional observer contract for:

• solve_linear_ode with dense, fixed-Krylov, and adaptive matrix-free exponential actions;
• DTMC/CTMC distribution forecasts;
• array and pytree probability/state representations;
• endpoint and requested-grid outputs where mathematically available.

An observer should produce a real-floating pytree, compose with jit/vmap/JVP/VJP, and avoid inflating the propagated state. Reuse existing deterministic interpolation only when its accuracy/order is defensible; otherwise evaluate at actual returned nodes.

Lazy residual/context split

Before adding matrix-free implicit ODE/DAE stages, define an internal split interface that can:

• evaluate residual only for root/Krylov/Jacobian-vector operations;
• evaluate (residual, algebraic_context) once when dynamics or observation needs the context;
• preserve the current full g → context → f derivative when context is consumed;
• keep failure-safe references and discrete-context float0 tangents correct;
• avoid tracing/materializing large context pytrees in residual-only products where possible.

This should remain an internal optimization unless a public API is demonstrably necessary.

Performance and AD requirements

• Explicitly disabled observers/context must preserve the minimal compiled path.
• Matrix-free implementations must not densify generators, Jacobians, or pytree states.
• Hand-coded JVP/VJP may reuse Krylov bases or factorizations where the primal is fixed and the math permits it.
• Test terminal and grid observations in primal/JVP/VJP, including multiple tangent/cotangent directions.
• Benchmark scalar/small dense systems and large matrix-free pytree systems on CPU/GPU.

Acceptance criteria

1. A written mathematical/API design covering observer timing, interpolation, and failure behavior.
2. Working observer output for linear exponential and fixed-chain forecast paths.
3. JVP/VJP and vmap tests for arrays and pytrees.
4. A residual-only/context-on-demand prototype for a matrix-free implicit path, or a documented decision to defer it with measured evidence.
5. No >5% compiled runtime regression in disabled array baselines.