Skip to content

chore(jax_assertions): assert sentinel default for use_jax_for_visualization#28

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/use-jax-for-vis-default
May 16, 2026
Merged

chore(jax_assertions): assert sentinel default for use_jax_for_visualization#28
Jammy2211 merged 1 commit into
mainfrom
feature/use-jax-for-vis-default

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

Summary

Phase 2 of the JAX visualization roadmap. PyAutoFit PR #1278 changes
Analysis(use_jax_for_visualization) from bool = False to
Optional[bool] = None — a sentinel that resolves to use_jax. The new
default-on behaviour for use_jax=True callers requires JAX at runtime to
exercise, so per the library convention (numpy-only unit tests; JAX-needing
assertions in autofit_workspace_test/scripts/jax_assertions/) the
sentinel-resolution checks live here, not in test_autofit/.

Scripts Changed

  • scripts/jax_assertions/fitness_dispatch.py — four new assertion functions
    (and the __main__ block) covering:
    • assert_use_jax_true_implicitly_turns_on_visualization — bare
      Analysis(use_jax=True) resolves _use_jax_for_visualization to True
    • assert_explicit_false_opts_out_when_use_jax_true — explicit
      use_jax_for_visualization=False overrides the sentinel
    • assert_explicit_none_resolves_to_use_jax — passing None is
      identical to omitting the argument
    • assert_use_jax_true_jit_dispatch_via_sentinel_default — end-to-end:
      Analysis(use_jax=True) with sentinel default actually wires
      _jitted_fit_from on first fit_for_visualization call

Upstream PR

PyAutoLabs/PyAutoFit#1278

Test Plan

  • python scripts/jax_assertions/fitness_dispatch.py — all 9
    assertions (5 pre-existing + 4 new) pass; the closing Dynesty
    af.Array fit completes and returns the expected shape.
  • Smoke tests pass for autofit_workspace_test.

🤖 Generated with Claude Code

…ization

Phase 2 of the JAX visualization roadmap. PyAutoFit PR #1278 flipped the
default of Analysis(use_jax_for_visualization) from False to None
(sentinel that follows use_jax). This workspace_test follow-up adds the
JAX-conditional sentinel-resolution assertions to scripts/jax_assertions/
fitness_dispatch.py, in line with the rule that library unit tests stay
numpy-only and JAX-needing assertions live here.

Four new assertions:
- assert_use_jax_true_implicitly_turns_on_visualization
- assert_explicit_false_opts_out_when_use_jax_true
- assert_explicit_none_resolves_to_use_jax
- assert_use_jax_true_jit_dispatch_via_sentinel_default

Linked to PyAutoFit PR #1278 and issue #1275.
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label May 16, 2026
@Jammy2211 Jammy2211 merged commit 79c3489 into main May 16, 2026
4 checks passed
@Jammy2211 Jammy2211 deleted the feature/use-jax-for-vis-default branch May 16, 2026 10:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant