test(jax): consolidate scripts/jax_assertions/ subdir for PyAutoFit#21
Merged
test(jax): consolidate scripts/jax_assertions/ subdir for PyAutoFit#21
Conversation
Reorganize jax assertions into a dedicated subdir mirroring the autolens_workspace_test/scripts/jax_assertions/ pattern. - git mv scripts/graphical/jax_assertions.py to scripts/jax_assertions/fitness_dispatch.py (preserves the 5 existing Fitness/JIT/pickle/visualization assertions added in 1caf713). - 3 new scripts holding assertions that previously lived as pytest tests in PyAutoFit/test_autofit/: - pytrees.py: 5 manual jax pytree roundtrip checks for Prior, Model, Collection, ModelInstance via the legacy direct-registration path. - enable_pytrees.py: 8 checks for the public autofit.jax.enable_pytrees / register_model API (model roundtrip, JIT compile, constant static-routing, kwarg-constant static, TuplePrior gradient flow, idempotency, Collection roundtrip). - nested.py: 8 checks comparing autofit.graphical.utils.nested_* utilities against jax.tree_util as reference (nested_get/set, ordering, items/paths, filter, map, update preserving NamedTuple). Each script ends with a `print("<name>: all assertions passed")` marker. All 4 scripts run successfully under PYAUTO_TEST_MODE-free conditions. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
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
Reorganize jax assertions into a dedicated
scripts/jax_assertions/subdir mirroringautolens_workspace_test/scripts/jax_assertions/. Adds 3 new scripts holding 20 assertions previously living as pytest tests inPyAutoFit/test_autofit/.Changes
git mv scripts/graphical/jax_assertions.py→scripts/jax_assertions/fitness_dispatch.py(preserves the 5 existing Fitness/JIT/pickle/visualization assertions added in1caf713)pytrees.pytest_autofit/jax/test_pytrees.pyjax.tree_util.register_pytree_node_classroundtrips for Prior, Model, Collection, ModelInstanceenable_pytrees.pytest_autofit/jax/test_enable_pytrees.pyautofit.jax.enable_pytrees/register_modelAPI: roundtrip, JIT compile, constant static-routing, kwarg static,TuplePriorgradient flow, idempotency, Collection roundtripnested.pytest_autofit/graphical/functionality/test_nested.pyautofit.graphical.utils.nested_*vsjax.tree_utilparity (get/set, ordering, items/paths, filter, map, update preserving NamedTuple)Companion PR
Pairs with PyAutoFit#new-pr, which deletes the source unit tests. Merge this PR first so the assertions exist on main before the source tests are removed.
Test plan
python scripts/jax_assertions/<name>.py(locally verified with thecuda_plugin_extension is not foundJAX warning being benign)fitness_dispatch.pyruns successfully (no behavior change from the path move)