revert: default use_jax_for_visualization to False (reverts #1278)#1280
Merged
Conversation
PR #1278 made `use_jax_for_visualization` default to follow `use_jax` (`Optional[bool] = None` resolving to `use_jax`). That caused every Nautilus quick-update under `use_jax=True` to evaluate `jax.jit(self.fit_from)(instance=instance)` where `instance` is a `ModelInstance` — a plain Python object that is not pytree-registered. JAX raises a `TypeError` trying to abstract it. On real pipeline runs (e.g. `z_projects/euclid/scripts/initial_lens_model.py`) the exception was swallowed deeper in the visualizer's outer guards; visible symptom was source-plane FITS files written all-zero and posteriors collapsing to the full prior on Einstein radius across every Euclid tile. Reverts the default to `bool = False`. Drops the sentinel-resolution block. Explicit opt-in (`use_jax_for_visualization=True`) and the existing `use_jax_for_visualization=True and not use_jax` warning both remain. No code in the PyAuto ecosystem relied on the implicit- on behaviour introduced by #1278. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2 tasks
Collaborator
Author
|
Workspace PR: PyAutoLabs/autofit_workspace_test#29 |
Jammy2211
pushed a commit
that referenced
this pull request
May 17, 2026
…inline) The Tests workflow has been red on `main` since 2026-05-16 09:48 UTC, the moment PR #1277 merged the `autofit[nss]` install extra. The 12 tests under `test_autofit/non_linear/search/nest/nss/` hit the script's own ImportError guard on `af.NSS()` because the CI install step only installs `[optional]`. A naïve fix (combine `[optional,nss]` into one pip install) hit a real dependency conflict — both extras pin `blackjax` but to different versions: - [optional] pins `blackjax>=1.2.0` (mainline, PyPI) - [nss] pins handley-lab fork @ ef45acd2 (~0.1.0b1.dev85+) The fork carries the `blackjax.ns.adaptive.init` entrypoint that mainline 1.2.x lacks, so it's not a "use the older one" merge — they're genuinely incompatible. pip rightly refuses `[optional,nss]` with `ResolutionImpossible`. Resolution: split into parallel jobs. - `unittest`: installs `[optional]` and runs the full test suite EXCLUDING `test_autofit/non_linear/search/nest/nss/`. This keeps `test_blackjax_nuts.py` (which requires mainline blackjax 1.2+) green. - `unittest_nss` (new): installs `[nss]` alone in a fresh env and runs ONLY the NSS test suite. Matrix is python 3.12 / 3.13 to match the main `unittest` job. Both jobs need to pass for the `Tests` workflow to be green. The existing `nss_install_smoke.yml` workflow stays as the weekly-cron upstream-drift canary on the [nss] git pins; this new job is the PR/push-time test gate. Fixes the failing Tests check observed on runs 25958838630 (#1277 merge), 25959492149 (#1278 merge), 25998822676 (#1280 revert merge), and 25999018708 (the first attempt at this fix that revealed the extras conflict). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
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
1 participant
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
Reverts PR #1278's change to
Analysis.__init__'suse_jax_for_visualizationdefault. The new default (Optional[bool] = None→ followuse_jax) caused every Nautilus quick_update underuse_jax=Trueto calljax.jit(self.fit_from)(instance=instance)whereinstanceis aModelInstance— a Python object that is not pytree-registered. JAX raises aTypeErrortrying to abstract it.On real pipeline runs (e.g.
z_projects/euclid/scripts/initial_lens_model.py) the exception was swallowed by the visualizer's outer guards; the visible symptom was source-plane FITS written all-zero and Einstein-radius posteriors collapsing to the full prior across every Euclid tile after the May 16 update. Reverting restores the previous behaviour: JIT visualization is opt-in only.Closes the regression introduced in #1278.
API Changes
Analysis.__init__(use_jax_for_visualization)default flips back fromOptional[bool] = None(followuse_jax) tobool = False. The sentinel-resolution block is dropped. Users wanting JIT visualization must passuse_jax_for_visualization=Trueexplicitly. The existinguse_jax_for_visualization=True and not use_jaxwarning remains. No public API surface was added — this is purely a default-value revert. See full details below.Test Plan
pytest test_autofit/analysis/test_use_jax_for_visualization.pypasses (6 passed, 1 skipped — verified locally)Analysis(use_jax=True)._use_jax_for_visualizationisFalse(default off)Analysis(use_jax=True, use_jax_for_visualization=True)._use_jax_for_visualizationisTrue(explicit opt-in still works)autofit_workspace_testPR ships next (workspace assertions updated to match new default —assert_use_jax_true_defaults_visualization_offreplaces the three sentinel assertions)Full API Changes (for automation & release notes)
Changed Signature
autofit.non_linear.analysis.analysis.Analysis.__init__(use_jax_for_visualization)— default flipped fromOptional[bool] = Nonetobool = False.Removed
Analysis.__init__:False.Migration
Analysis(use_jax=True)→ JAX visualization stays off. No code change required.use_jax_for_visualization=Trueexplicitly. As of this revert, no code in the PyAuto ecosystem relies on the implicit-on behaviour.🤖 Generated with Claude Code