Planning Note: JAX Backend, Jacobians, and Optimizer

Status (2026-07-11): Phases A–D of this plan are implemented (eval_jax.py, can_lower_jax_2d, fit_model_jax + Dfun plumbing); Phase E (optimizer replacement) is deferred for lack of evidence that lmfit is the bottleneck. This note is preserved as the planning-time rationale.

Summary

The current NumPy GIR backend is ready to be treated as complete for the current lowered-evaluator scope. A JAX port should proceed as a new track, not as unfinished work that must land before GIR / eval / scheduler can be called “ready”.

That conclusion rests on the current architecture:

• evaluate_2d(plan, theta) and evaluate_1d(plan, theta) already expose theta-level evaluator entry points.

• ScheduledPlan2D and ScheduledPlan1D are already mostly packed-array execution plans rather than live model objects.

• fit_model_gir already extracts the optimizer-visible theta vector from the full parameter list and routes into the compiled evaluator.

• Parity coverage already exists for static 1D/2D models, profiles, convolution/IRF, and subcycle-aware cases in tests/test_evaluate_1d.py, tests/test_evaluate_2d.py, and tests/test_gir_integration.py.

What remains is real backend work, not “finish the GIR architecture” work.

Current readiness

The current implementation already has several properties that make a JAX track feasible:

• The 2D evaluator is already organized as a pure-looking function of (plan, theta) -> spectrum, even though the current implementation uses NumPy scratch buffers and mutation internally.

• The hot loops are plan-driven, not theta-driven. In practice this means the branch structure is fixed once schedule_2d(graph) has run, which is a good fit for a compiled backend.

• The scheduler has already done the hard semantic work: expression binding, dynamics grouping, convolution lowering, profile lowering, and subcycle lowering all happen before the evaluator runs.

• The current lmfit boundary is narrow. residual_fun calls a fit function, and fit_model_gir already converts the full parameter vector into the compact optimizer vector expected by the evaluator.

This is enough to start a JAX branch immediately.

It is not necessary to keep the current GIR branch open for additional “JAX readiness” cleanup before declaring the NumPy backend complete.

What still blocks a full JAX path

1. Backend-agnostic prep

These items are not blockers for calling the NumPy GIR work done, but they are useful cleanup before or during a JAX port:

• Flatten expression storage. ScheduledPlan2D.expr_programs and ScheduledPlan2D.profile_expr_programs are still Python lists of ExprProgram objects. They should move to fully packed arrays (for example CSR-style instruction storage plus per-program offsets).

• Keep schedule data fully array-native. The main plan is already close to this ideal; the remaining goal is to avoid Python-wrapped structures in the hot path altogether.

• Preserve a stable theta contract. The evaluator boundary evaluate(plan, theta) is already the right API. New work should protect that contract rather than pushing JAX concerns back into model objects or fit-time parsing.

• Kernel-matrix convolution has landed (2026-07): the lowered convolution path now evaluates registry kernels elementwise on a precomputed dt matrix and applies a quadrature-weighted matmul (kernel-matrix-convolution.md). All array shapes in the convolution path are theta-independent, so the former jit blocker (per-theta kernel support lengths) is gone.

Of these, expression flattening is the one prep item most worth doing even if the eventual backend choice changes again.

2. JAX evaluator port

The main technical work is in the evaluator itself:

• Mutation-heavy NumPy code must become functional JAX code. The current 2D evaluator uses .copy(), slice assignment, and in-place accumulation over scratch arrays. A jitted JAX version will need explicit functional updates or loop-carried state.

• Python callable dispatch must become JAX-native dispatch. The current evaluator relies on Python dispatch tables (OP_DISPATCH, DYNAMICS_DISPATCH, PROFILE_DISPATCH). A jitted backend will need either explicit JAX control flow or trace-time unrolling over the scheduled ops.

• Host-side checks must stay outside the jitted region. Shape checks and Python exceptions are fine at the outer boundary, but should not live inside the compiled path.

• SciPy-dependent kernels need JAX-compatible replacements. The kernel-matrix change (kernel-matrix-convolution.md) retired the SciPy convolution utilities from the lowered path, and the removal of voigtCONV/lorentzCONV (2026-07) retired wofz from functions/time.py entirely — the conv path is now JAX-expressible (erfc for the Gaussian edge masses exists in jax.scipy.special). The energy-domain Voigt profile still uses wofz and remains the one special-function gap.

3. Jacobian and optimizer work

Even after a working JAX evaluator exists, Jacobian / optimizer work is still its own layer:

• fitlib.py currently constructs lmfit.Minimizer without a Jacobian hook. If we want analytic Jacobians while keeping lmfit, we need explicit Dfun plumbing.

• A fully custom JAX optimizer is a larger decision than “use JAX for derivatives”. It also means deciding how to replace or re-scope:

  • existing result objects and reporting,

  • two-stage fitting workflow,

  • confidence-interval tooling,

  • MCMC integration,

  • parameter bounds / transformations,

  • writeback into the current model/result surface.

Because of that, “JAX evaluator + analytic Jacobian” and “replace lmfit” should be treated as separate milestones, not one first step.

Recommendation on optimizer strategy

The recommended first optimizer milestone is:

1. build a JAX evaluator,

2. derive a Jacobian from it,

3. keep lmfit as the outer optimizer at first.

Reasons:

• It isolates the performance question. We learn whether JAX + Jacobian is actually worthwhile before committing to a larger optimizer rewrite.

• It preserves the existing user-facing fit workflow while the new backend proves itself.

• It avoids bundling backend correctness risk with optimizer-behavior risk in the same first implementation.

A custom JAX optimizer may still be the right long-term direction, but it should come after we have:

• a parity-validated JAX evaluator,

• a working analytic Jacobian path,

• and concrete evidence that lmfit itself is the next real bottleneck.

Recommended implementation order

Phase A: backend-agnostic prep

• Flatten expr_programs and profile_expr_programs into packed arrays.

• Keep the evaluator API unchanged: evaluate(plan, theta) -> ndarray.

• Add any small helper refactors that make the evaluator easier to port without changing behavior.

Phase B: first JAX evaluator slice

Target file-level 2D fits first; that is where the payoff is.

Initial scope:

• static component ops,

• dynamics groups,

• arithmetic expressions,

• no profiles,

• no convolution,

• no Voigt.

The goal of this phase is to validate the JAX evaluator structure, not to land full coverage immediately.

Phase C: widen JAX coverage to the existing GIR surface

Add the remaining lowered features incrementally:

• profile-varying parameters,

• subcycle-aware dynamics,

• resolved-trace convolution (kernel-matrix form; see kernel-matrix-convolution.md),

• Voigt / special-function support.

Each widening step should ship with direct parity tests against the existing NumPy GIR evaluator, not only against MCP.

Phase D: analytic Jacobian

• Define a residual function at the theta level.

• Differentiate that residual or model output with JAX.

• Wire the resulting Jacobian into the current fit flow while keeping the rest of the API stable.

At this point we can benchmark fit-level speedup rather than only per-call evaluator speedup.

Phase E: optional JAX-native optimizer

Only after Phase D should we decide whether replacing lmfit is justified.

This phase is explicitly optional. It should be driven by measured bottlenecks, not by the assumption that “using JAX” automatically implies “replace the whole optimizer stack”.

What does not belong in the closing GIR branch

The following should be treated as next-track work, not as unfinished business for the current lowered-evaluator branch:

• JAX backend porting work

• Jacobian plumbing

• optimizer replacement

• project-level fit lowering

• mixed-backend execution

Those are legitimate follow-ons, but they are not prerequisites for calling the current NumPy GIR evaluator complete.

Success criteria for the JAX track

A JAX branch should be considered successful when all of the following are true:

• A JAX evaluator matches the NumPy GIR evaluator to practical floating-point tolerance on the supported subset.

• The widened JAX path covers the same intended feature surface as the current NumPy GIR backend, or clearly documents any temporary exclusions.

• Fit-level benchmarks show a meaningful win from JAX alone or from JAX + analytic Jacobian.

• The existing NumPy GIR path remains available as the conservative reference implementation until the JAX path has earned default status.

Bottom line

The current scheduler / evaluator design is already good enough to start a JAX project now.

There is still meaningful implementation work ahead, but it belongs on a new backend branch. It should not be framed as “we cannot call GIR ready until we do this.” The right framing is: the NumPy GIR path is the completed reference backend, and JAX is the next backend experiment built on top of it.