Repository navigation
feat: lazily JIT-cache fit_for_visualization when flag set - #1229
Merged
Merged
Conversation
Replace the pass-through stub in Analysis.fit_for_visualization with a lazy jax.jit cache on the instance. With use_jax_for_visualization=True, the first call compiles self.fit_from and stores it as _jitted_fit_from; subsequent calls reuse it. Flag defaults to False so no existing fit is affected until its Fit* return type has been pytree-audited. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Merged
2 tasks done
Collaborator
Author
|
Workspace PR: PyAutoLabs/autolens_workspace_test#33 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Replace the pass-through stub in
Analysis.fit_for_visualizationwith a lazyjax.jitcache. Withuse_jax_for_visualization=True, the first call compilesself.fit_fromand stores it on the instance as_jitted_fit_from; subsequent calls reuse the compiled function. Default behaviour (flag off) is unchanged — this is a controlled opt-in ahead of the remainingFit*pytree registrations.API Changes
Analysis.fit_for_visualizationnow lazily constructs and cachesjax.jit(self.fit_from)on the instance whenuse_jax_for_visualization=True. Flag-off behaviour and all public signatures are unchanged. See full details below.Test Plan
test_autofit/analysis/test_use_jax_for_visualization.pyupdated to exercise both the JIT path (asserts_jitted_fit_fromis created on first call and reused on second) and the no-flag pathtest_autofit/suite passes (1232 tests)modeling_visualization_jit.pyin the linked workspace PR: 7.6× speedup compile → cached, Nautilus live-search fires the JIT path and writessubplot_fit.pngFull API Changes (for automation & release notes)
Changed Behaviour
Analysis.fit_for_visualization(instance)— whenuse_jax_for_visualization=Trueon the instance, lazily constructsjax.jit(self.fit_from)on the first call, caches it asself._jitted_fit_from, and returns the compiled function's output for every subsequent call. Falls back to plainself.fit_from(instance)when the flag is off.Added (internal)
Analysis._jitted_fit_from— lazy-initialized attribute;Noneuntil the firstfit_for_visualizationcall with the flag on.Migration
No migration needed for flag-off callers (the default). Callers that set
use_jax_for_visualization=Truenow get a JIT-compiledfit_from— theirFit*return type and all nested types must be pytree-registered or the compile will error. Seeadmin_jammy/prompt/autolens/fit_*_pytree_*.mdfor the variant matrix.🤖 Generated with Claude Code