perf(galaxy): skip the second deflection call for the over-sampled grid at sub-size 1 (PyAutoArray#514) - #595
Merged
Merged
Conversation
…id at sub-size 1 (PyAutoArray#514) `Galaxy.traced_grid_2d_from` called `deflections_yx_2d_from` twice — once for the slim grid and once for `grid.over_sampled`. When the grid's over-sampler is uniform at sub-size 1 — the pixelization grid in every CPU likelihood cell — `grid.over_sampled` is the slim grid in the same order, so the second call recomputes a bit-identical answer at full cost and is now reused instead. The sub-size is host numpy, so the guard is a static Python bool and the branch is resolved at JAX trace time rather than becoming a traced `cond`. The equivalent guard in `Tracer.traced_grid_2d_list_from` (PyAutoLens) measures the tracer/raw deflection ratio on hst dropping from 2.52x to 1.41x for Isothermal and 1.94x to 0.88x for PowerLaw, with every pinned value still PASSING — the change is bit-identical. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HWjPT94MPbEHT45kJmDpDh
This was referenced Sep 2, 2026
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
Phase 1 of the
numpy-deflections-cpuepic (PyAutoArray#514).Galaxy.traced_grid_2d_fromcalleddeflections_yx_2d_fromtwice on aGrid2D: once for the grid and once forgrid.over_sampled. At a uniform over-sample size of 1 the over-sampled grid is the slim grid in the same order, so the second call doubled the deflection cost for a bit-identical result. Same guard as the companion PyAutoLens tracer change: the over-sample size is host numpy, so the branch is static at JAX trace time; sub-size > 1 is unchanged.API Changes
None — internal changes only.
See full details below.
Test Plan
test_autogalaxy: 1158 passed (new tests intest_autogalaxy/galaxy/test_galaxy.py: traced.over_sampledequals the separately traced over-sampled grid and differs from the untraced one, at sub-size 1 and 4; at sub-size 1 it equals the traced slim grid exactly).autolens_profilinglens cells hst: deflection pins PASSED (rtol 1e-6).Full API Changes (for automation & release notes)
Changed Behaviour
autogalaxy.galaxy.galaxy.Galaxy.traced_grid_2d_from— when every entry ofgrid.over_sample_sizeis 1, the over-sampled traced grid is aGrid2DIrregularcopy of the traced slim grid instead of a seconddeflections_yx_2d_fromcall (bit-identical values).Generated by the PyAutoLabs agent workflow. Companion PRs: PyAutoArray (
to_grid/Grid2D.over_sampled), PyAutoLens (Tracer.traced_grid_2d_list_from), autolens_profiling (scripts/lens/measurement package).Companion: PyAutoLabs/PyAutoArray#516