Skip to content

JAX evaluator backend with analytic Jacobians (v0.12.0) - #37

Merged
InfinityMonkeyAtWork merged 12 commits into
mainfrom
jax-backend
Jul 12, 2026
Merged

JAX evaluator backend with analytic Jacobians (v0.12.0)#37
InfinityMonkeyAtWork merged 12 commits into
mainfrom
jax-backend

Conversation

@InfinityMonkeyAtWork

Copy link
Copy Markdown
Owner

Adds an experimental JAX backend for 2D fits, opt-in via Project.spec_fun_str =
"fit_model_jax" and the [jax] extra. Implements Phases A–D of
docs/design/jax-planning.md; execution record in docs/design/archive/jax-backend.md.

  • Phase A (backend-agnostic): expression programs flattened to packed CSR arrays; plans
    fully array-native, evaluator API unchanged.
  • Phases B+C: jitted 2D evaluator covering the full can_lower_2d surface (dynamics,
    expressions, subcycles, profiles, kernel-matrix convolution, Voigt via Weideman wofz,
    <1e-12 vs scipy). Gate composes with can_lower_2d, so the fallback chain is JAX → NumPy
    GIR → MCP by construction.
  • Phase D: jax.jacfwd analytic Jacobian wired into lmfit leastsq via Dfun. Voigt 2D fit:
    25.3 s → 7.1 s (~3.6x incl. compiles), identical redchi, nfev 8 → 2.
  • Testing: every comparison in test_evaluate_2d.py now asserts three-way parity
    (interpreter/NumPy/JAX) at 1e-12; JAX-specific tests cover gate, Jacobian-vs-FD, wofz,
    e2e fit. Main CI installs .[dev,jax]; min-versions stays jax-free (jax needs numpy≥2.0).
  • Known limitations (documented): no runtime value checks on traced params (LinBack
    ordering, kernel positivity), no parallel-worker MCMC, imports enable global x64, 1D fits
    stay on the NumPy path.
  • Phase E (optimizer replacement) deferred with rationale; follow-on design notes in
    docs/design/ui.md and docs/design/project-level-fits.md.

@InfinityMonkeyAtWork
InfinityMonkeyAtWork merged commit 0b2c2ca into main Jul 12, 2026
4 checks passed
@InfinityMonkeyAtWork
InfinityMonkeyAtWork deleted the jax-backend branch July 12, 2026 16:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant