Skip to content

Project-level shared fits on the fused JAX backend (v0.13.0) - #39

Merged
InfinityMonkeyAtWork merged 7 commits into
mainfrom
jax-project-fit
Jul 14, 2026
Merged

Project-level shared fits on the fused JAX backend (v0.13.0)#39
InfinityMonkeyAtWork merged 7 commits into
mainfrom
jax-project-fit

Conversation

@InfinityMonkeyAtWork

Copy link
Copy Markdown
Owner

Summary

Adds a compiled fast path for project-level shared-parameter fits: when
Project.spec_fun_str == "fit_model_jax" and every file's model passes the
can_lower_2d + can_lower_jax_2d gates, Project.fit_2d() runs one fused,
jitted JAX program over all files (per-file plan trace, static fit-window
slices, flattened + concatenated residuals) with an analytic joint Jacobian
via jacfwd wired into lmfit's Dfun.

  • New spectra.pack_project_theta (combined-theta gather arrays) and
    eval_jax.make_project_evaluator_2d_jax / make_project_jacobian_2d_jax.
  • Whole-project fallback to fit_project_mcp if any file misses a gate or
    jax isn't installed — no mixed backends; the backend in use is reported in
    the fit output.
  • Heterogeneous energy/time grids across files are supported (per-file
    static window slices, concatenated residual vector).
  • Measured speedups vs the interpreter: 2.6x (2 files) → 7.6x (4) → 17x (8);
    XLA compile stays sub-linear (< ~0.4 s at 8 files).

Tests

  • TestProjectFitJax: interpreter/JAX end-to-end parity, truth recovery,
    heterogeneous grids, fallback gating (non-lowerable file, jax missing).
  • Full suite: 978 passed; sphinx -W docs build clean.

Docs / release

  • Design note docs/design/project-level-fits.md updated with implementation
    status and measured results; supported_models.md and
    roundtrip_test_matrix.md updated for the new backend axis.
  • CHANGELOG.md 0.13.0 entry; version bumped to 0.13.0. Tag v0.13.0 after
    merge to trigger the PyPI release workflow.

@InfinityMonkeyAtWork
InfinityMonkeyAtWork merged commit d0d53c0 into main Jul 14, 2026
4 checks passed
@InfinityMonkeyAtWork
InfinityMonkeyAtWork deleted the jax-project-fit branch July 14, 2026 16:15
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