ADR-0004: Framework-agnostic ML surrogate wrapper (forward + adjoint, PyTorch & JAX, ONNX)

Field	Value
Kind	ADR
Status	Accepted
Decided	2026-06
Deciders	Antoine Berchet
Supersedes	—
Superseded by	—

Context

CIF plugs ML surrogates into a variational inversion framework, which needs not just a forward operator but its adjoint for gradients. Binding CIF to a single autodiff ecosystem would couple the framework to one vendor’s lifecycle and exclude models built in the other.

Decision

The ML wrapper is framework-agnostic: it exposes forward and adjoint operators in both PyTorch and JAX, with ONNX export for deployment / inference.

Consequences

• Surrogates slot into the existing variational machinery because the adjoint is a first-class part of the operator contract.

• ONNX export decouples the training environment from inference / deployment.

• Costs: two backend code paths to maintain, and ONNX operator-coverage limits can constrain which models are exportable.

• Adjoint correctness must be verified against the forward via gradient checks — a required validation step for any surrogate, not an optional one.