bug: jit/imaging/pixelization.py — JIT vs eager mapping matrix shape mismatch

[Jammy2211]

Summary

On clean main (autolens_workspace_developer@4e99bf9, PyAutoLens 2026.5.14.2), jax_profiling/jit/imaging/pixelization.py crashes at the regularized-reconstruction step with a JAX broadcasting error. The JIT-traced blurred_mapping_matrix has shape (15361, 1285) whereas the eager path produces (15361, 1225), so the downstream curvature matrix (1285×1285) cannot be added to the regularization matrix (1225×1225).

Surfaced by the autolens_profiling Phase 1 follow-up smoke (issue #67); split out of that triage because it's a pre-existing JIT-vs-eager bug independent of the point_source regression drift.

Reproduction

cd autolens_workspace_developer
source ../activate.sh  # or set NUMBA_CACHE_DIR / MPLCONFIGDIR / PYTHONPATH manually
python jax_profiling/jit/imaging/pixelization.py

Output (relevant tail)

  blurred_mapping_matrix (JIT) shape: (15361, 1285)
  blurred_mapping_matrix shape: (15361, 1225)
...
  curvature_matrix shape: (1285, 1285)
...
  regularization_matrix shape: (1225, 1225)

--- Step 12: Regularized reconstruction ---
Traceback (most recent call last):
  File ".../jax_profiling/jit/imaging/pixelization.py", line 660, in <module>
    reconstruction = compute_reconstruction(
  File ".../jax_profiling/jit/imaging/pixelization.py", line 652, in compute_reconstruction
    curvature_reg_matrix = curvature_matrix + regularization_matrix
TypeError: add got incompatible shapes for broadcasting: (1285, 1285), (1225, 1225).

Diagnosis hint

• 1225 = 35×35 (likely the rectangular pixelization native mesh shape, 35x35 pixels).
• 1285 - 1225 = 60. Possibly 60 extra source pixels from over-sampling / mesh expansion in the JIT path that the eager path doesn't apply.
• The blurred_mapping_matrix is being computed with a different mesh shape on the JIT path vs the eager path, even though both should be reading the same Pixelization model.

Likely culprits to investigate (somewhere in PyAutoArray / PyAutoGalaxy):

• Mapper-side mesh construction under xp=jnp.
• Pixelization source-mesh shape derivation when xp=jnp vs xp=np.
• Over-sampling / image-mesh padding applied conditionally.

Context

• All other JIT scripts on clean main behave consistently:
  • imaging/mge.py, interferometer/{mge,pixelization,delaunay}.py, datacube/delaunay.py — PASS.
  • imaging/delaunay.py — crashes with log_evidence -inf (separate issue).
  • point_source/{image_plane,source_plane}.py — drift in regression constants (triaged in triage: drift in jit/ regression constants (autolens_profiling F1) #67).

This is the only pixelization (rectangular) JIT script that crashes — interferometer/pixelization works, which is suspicious enough to point at imaging-specific mapper code.

Related

• Triage parent: triage: drift in jit/ regression constants (autolens_profiling F1) #67 (autolens_profiling F1 regression-constant drift)
• Phase 1 mirror PR: feat: Phase 1 — mirror JIT likelihood profiling scripts autolens_profiling#2

[Jammy2211]

Today's diagnostic update (2026-05-16)

Smoking through the full profile suite for Phase 5 validation, I confirmed this bug is still firing on PyAutoLens 2026.5.14.2. Two of the other three known bugs in the autolens_profiling smoke suite turned out to be script-level bugs (#69 was rank-deficient regularisation handling in the rebuild formula; #514 was stale prior means relative to the simulator). But #68 looks like a real library issue, not a script bug — flagging the diagnostic so a maintainer with PyAutoArray expertise can dig in.

The two paths that disagree

likelihood/imaging/pixelization.py constructs blurred_mapping_matrix (= operated_mapping_matrix) three times, and two of them disagree on shape:

Path	Where	Output shape
A (eager from FitImaging)	inversion.operated_mapping_matrix (line 507)	(15361, 1225)
B (eager from dataset.psf)	dataset.psf.convolved_mapping_matrix_from(mapping_matrix=inv_mapper.mapping_matrix, mask=dataset.mask, xp=jnp) (line 539)	(15361, 1225)
C (JIT, full rebuild)	al.FitImaging(..., xp=jnp).inversion.operated_mapping_matrix inside JIT closure (line 557)	(15361, 1285)

Paths A and B agree at 1225. Path C diverges at 1285. The mismatch is 60 columns — exactly suggestive of border-relocator behaviour differing between the eager FitImaging (path A, which built inv_mapper once with use_border_relocator=True) and the JIT-rebuilt FitImaging (path C, which calls al.FitImaging(..., settings=al.Settings(use_border_relocator=True), xp=jnp)).

Likely candidates:

1. Border relocator under xp=jnp drops fewer (or more) source-mesh pixels than the numpy path. The 60-pixel difference corresponds to about an edge-ring at the source mesh.
2. Mesh / overlay grid construction does something different in the JIT path — e.g. the Overlay image_mesh grid hits a different rounding or different image_plane_pixelization_grid size when traced under xp=jnp.
3. Some preallocation in the JIT path keeps "zero" columns (e.g. for relocated pixels) where the eager path strips them.

Reproduction

cd autolens_profiling  # at PyAutoLabs/autolens_profiling main
python likelihood/imaging/pixelization.py

Crashes at Step 12 (regularised reconstruction) with:

TypeError: add got incompatible shapes for broadcasting: (1285, 1285), (1225, 1225).

— because the curvature matrix is built from the JIT-shape mapping matrix (1285) but the regularisation matrix is built from the eager-shape inversion (1225).

Affected scripts in the mirror

The same shape divergence applies to autolens_profiling/likelihood/imaging/pixelization.py (Phase 1 mirror of this _developer script). When #68 is fixed upstream, the mirror picks it up automatically since it's identical code.

Refs

• autolens_profiling#1 (Phase 1 mirror parent) — surfaces this bug on the mirror side via the same crash.
• Sibling fix-on-script bugs from today's smoke:
  • fix: imaging/delaunay log_evidence -inf via rank-stripped regularisation (#69) autolens_profiling#13 (fix for #69 — script-level)
  • fix: point_source priors aligned to simulator truth + refresh constants (#514) autolens_profiling#14 (fix for #514 — script-level)

These two were closable as script fixes; this one (#68) really does need PyAutoArray-side investigation, hence the more careful diagnostic write-up.