Repository navigation
feat(regularization): numba kernels for the numpy split-regularization assembly (#267 lever 1) - #553
Merged
Merged
Conversation
…n assembly Replace the two pure-Python loops behind reg_split_np_from and pixel_splitted_regularization_matrix_np_from with @numba_util.jit kernels that reproduce the reference accumulation order bit-for-bit. The Python bodies are retained as _reg_split_reference and _pixel_splitted_regularization_matrix_reference and pinned against the kernels with np.array_equal in test_regularization_util_numba.py. Behaviour changes (deliberate, recorded): reg_split_np_from copies mappings and sizes on entry instead of mutating the interpolators' cached stencil tables; an empty row inserts its centre pixel at column 0 instead of a stale column (the size == 0 j-leak); a full stencil with no self pixel raises exc.InversionException instead of writing past the table width. JAX branch untouched (PyAutoArray#536 owns it). Local single-thread timing at P=1500, K=4: reference 106.3 ms -> kernel 1.23 ms. Lever 1 of PyAutoLabs/autolens_profiling#267. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Collaborator
Author
|
Paired measurement PR (cells, JSONs, verdict note): PyAutoLabs/autolens_profiling#269 — RAL jobs 343345 (CPU A/B, 1.38x) and 343346 (A100, identity). |
This was referenced Sep 16, 2026
Merged
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
Lever 1 of PyAutoLabs/autolens_profiling#267 (phase 3 of the
fixed-lens-light-numba-cpuepic): the numpy/numba split-regularization assembly is now two@numba_util.jit()kernels instead of two pure-Python loops.On the production numba CPU path at HST Delaunay N=1500,
inversion.regularization_matrixwas the single largest site in the source-only likelihood call: 111.5 ms of 404.6 ms (27.6 %), invariant under the NNLS memo.regularization_util.pycontained zero numba; PyAutoArray#536 optimised the JAX branch only.What changed
reg_split_np_fromandpixel_splitted_regularization_matrix_np_fromkeep their names and signatures and now call_reg_split_kernel/_pixel_splitted_regularization_matrix_kernel, which reproduce the original accumulation order bit-for-bit (the2e-8diagonal first, both-triangle adds inl/l+morder, the trailing/= 2.0last;fastmath=False)._reg_split_referenceand_pixel_splitted_regularization_matrix_reference. They remain the ground truth for the JAX branch tests (test_pixel_splitted_jax.py) and are now also pinned against the kernels withnp.array_equalin the newtest_regularization_util_numba.py(8 tests: narrow/wide DelaunayNN-shaped tables, K-NN tables throughreg_split, no input mutation, thesize == 0row, the overflow raise, and an end-to-endConstantSplit.regularization_matrix_fromagainst the reference chain).pixel_splitted_regularization_matrix_from/reg_split_fromis byte-identical tomain.API Changes
Three deliberate behaviour changes on the numpy path, none of which affects any caller in PyAutoArray, PyAutoGalaxy, PyAutoLens, PyAutoCTI or PyAutoFit (all consume the returned tuple):
reg_split_np_fromno longer mutates its arguments. It copiessplitted_mappingsandsplitted_sizeson entry and returns the completed tables; the interpolators'cached_propertystencil tables are no longer rewritten by the first regularization-matrix call. The returned mappings keep their full width including the reserved pad column, so the(mappings, sizes, weights)contract is unchanged. Side effect:test_adapt_power_jax.pynow compares its numpy and JAX legs from the same inputs (previously the JAX leg ran on already-mutated tables).size == 0) inserts its centre pixel at column 0. The old loop leftjat the previous row's value and wrote at a stalej + 1. The kernel uses an explicitinsert = sizes[i].exc.InversionExceptionnaming the row and table width, instead of anIndexError(Python) or silent out-of-bounds write (numba has no bounds checking).Measurements
Local single-thread sanity (P=1500, K=4, warm, median of 20): reference 106.3 ms → kernel 1.23 ms (86x);
reg_split9.13 → 0.040 ms.RAL A/B (
gpupartition CPUs-only, 1 thread numba + BLAS, private merge-base control5e2bc0f4vs feature onPYTHONPATH, phase-2 Leg D cell: route b, memo ON):inversion.regularization_matrix(ms)5e2bc0f4bae9296e1.38x on the whole call (1.383x on medians); the H site falls from 26.8 % to 2.0 % of the call. Both arms ran on one node back-to-back; the feature arm entered at a higher load average than the control, so 1.38x is a lower bound. The control arm reproduces phase 2's Leg D (404.6 ms / 111.5 ms) to +2.2 % / −0.8 %.
Witness inside the feature arm on RAL (
fixed_light_numba_levers_witness.py): the library'sinversion.regularization_matrixisarray_equalto the retained reference chain (max |diff| 0.0); the interpolator's cached stencil tables are byte-identical after the inversion; the fit's log evidence with the kernels vs the monkeypatched references differs by exactly 0.0 (gate ≤ 1e-9).A100 row (job 343346,
delaunay.pybreakdown cell, fp64,AdaptSplit): identity, as expected. The GPU cell dispatches to the JAX branch, so the numpy kernels are never reached there. The H prefix agrees to 0.03 % (7.3318 vs 7.3298 ms) and the whole call to 0.04 %;log_evidence_termsare bit-identical between arms. Finding for the record: withInterpolatorDelaunaythe stencil table is K = 4, soSPLIT_REG_COMPACT_WIDTH = 12/SPLIT_REG_WIDE_ROW_BUDGET = 256are inert at HST Delaunay N=1500 and can only be tuned on a cell with K > 12.Full tables, gate roll-up and provenance land in the paired autolens_profiling PR (
results/notes/).Testing
pytest test_autoarray/inversion/regularizations -q— 118 passedpytest test_autoarray/inversion/pixelization/interpolator -q— 58 passedpytest test_autoarray/inversion/inversion -q -x— 212 passedblack --checkcleanCloses nothing on its own; tracked by PyAutoLabs/autolens_profiling#267 (lever 1 of 3).
Heart gate
pyauto-heart readinesswas RED (score 20) at ship time with two release-chain reasons unrelated to this branch — "install verification FAILED (testpypi; checks F)" and "release validation FAILED (stage integrate)" — acknowledged by the human on 2026-09-16 for PR-open only. Merge stays a human act via/prm.🤖 Generated with Claude Code