perf(tracer): skip the second ray-trace of the over-sampled grid at sub-size 1 (PyAutoArray#514) - #718
Merged
Conversation
…ub-size 1 (PyAutoArray#514) `Tracer.traced_grid_2d_list_from` traced the input `Grid2D` and then traced `grid.over_sampled` a second time to fill the returned grids' `over_sampled` attribute. 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 trace recomputes a bit-identical answer at full cost. 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`. Measured on `autolens_profiling/scripts/lens/deflections/total.py --instrument hst` (the tracer/raw ratio, which normalises out laptop variance): Isothermal 2.52x -> 1.41x, PowerLaw 1.94x -> 0.88x, IsothermalSph 2.78x -> 1.61x, PowerLawSph 2.43x -> 1.38x. Every pinned value still PASSES — the change is bit-identical. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HWjPT94MPbEHT45kJmDpDh
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
Phase 1 of the
numpy-deflections-cpuepic (PyAutoArray#514).Tracer.traced_grid_2d_list_fromtraced aGrid2Dtwice: once for the grid and once forgrid.over_sampled, unconditionally. At a uniform over-sample size of 1 (the pixelization grid of every CPU likelihood evaluation,over_sample_size_pixelization=1) the over-sampled grid is the slim grid in the same order, so the second trace doubled every deflection-angle calculation for a bit-identical result. The over-sample size is host numpy, so the new guard is a static Python bool at JAX trace time, not a tracedcond; at sub-size > 1 the existing second trace runs unchanged.Measured with
autolens_profiling/scripts/lens/deflections/total.py --instrument hst(15,361 pixels,OMP_NUM_THREADS=1): tracer-to-raw-call ratio Isothermal 2.52× → 1.41×, PowerLaw 1.94× → 0.88× (the raw call includes the direct-Grid2Ddecorator overhead the companion PyAutoArray PR removes). Every deflection pin passes at rtol 1e-6; thepixelization_numba.pyhst likelihood pin (27661.910133664103, rtol 1e-6) passes. numpy and JAX (jax.jit) results agree at sub-size 1 and 4.API Changes
None — internal changes only. The traced
Grid2D.over_sampledat sub-size 1 is now aGrid2DIrregularcopy of the traced slim grid rather than a second, identical trace.See full details below.
Test Plan
test_autolens: 576 passed (new tests intest_autolens/lens/test_tracer.py: traced.over_sampledequalstraced(grid.over_sampled)and differs from the untraced grid, at sub-size 1 and 4; at sub-size 1 it equals the traced slim grid exactly).jax.jittraced grids agree at sub-size 1 and 4.autolens_profilinglens cells +pixelization_numba.pyhst: pins PASSED.Full API Changes (for automation & release notes)
Changed Behaviour
autolens.lens.tracer.Tracer.traced_grid_2d_list_from— when every entry ofgrid.over_sample_sizeis 1, the over-sampled traced grids are built from the already-traced slim grids instead of a second ray-trace ofgrid.over_sampled(bit-identical values; one fewer duplicate subgraph underjax.jit).Generated by the PyAutoLabs agent workflow. Companion PRs: PyAutoArray (
to_grid/Grid2D.over_sampled), PyAutoGalaxy (Galaxy.traced_grid_2d_from, same guard), autolens_profiling (scripts/lens/measurement package).Companion: PyAutoLabs/PyAutoArray#516
Companion: PyAutoLabs/PyAutoGalaxy#595