diff --git a/.gitignore b/.gitignore
index 5923432..8ed7f40 100644
--- a/.gitignore
+++ b/.gitignore
@@ -23,3 +23,26 @@ __jax_cache__/
 # Local virtualenvs
 .venv/
 venv/
+
+# Auto-simulated input datasets — regenerated by profiler scripts, not committed.
+# The pre-existing datasets (e.g. imaging/hst, interferometer/sma, point_source/simple/*.json)
+# remain tracked and are NOT covered by these rules.
+dataset/cluster/
+dataset/group/
+dataset/multi/
+dataset/imaging/simple/
+dataset/interferometer/simple/
+dataset/point_source/simple/data.fits
+dataset/point_source/simple/noise_map.fits
+dataset/point_source/simple/psf.fits
+dataset/point_source/simple/point_dataset_positions_only.csv
+
+# Datacube simulator outputs — regenerated each run
+dataset/datacube/sim_*/
+
+# Profiling result artifacts — written by scripts under likelihood/, simulators/, searches/.
+# Each run produces fresh versioned JSON+PNG; Phase 4's dashboard will explicitly select
+# which artifacts to commit. Until then, results stay local.
+results/likelihood/
+results/simulators/
+results/searches/
diff --git a/dataset/imaging/hst/data.fits b/dataset/imaging/hst/data.fits
new file mode 100644
index 0000000..f0f9c28
Binary files /dev/null and b/dataset/imaging/hst/data.fits differ
diff --git a/dataset/imaging/hst/noise_map.fits b/dataset/imaging/hst/noise_map.fits
new file mode 100644
index 0000000..fee51bf
Binary files /dev/null and b/dataset/imaging/hst/noise_map.fits differ
diff --git a/dataset/imaging/hst/psf.fits b/dataset/imaging/hst/psf.fits
new file mode 100644
index 0000000..1f7babf
Binary files /dev/null and b/dataset/imaging/hst/psf.fits differ
diff --git a/dataset/imaging/hst/tracer.json b/dataset/imaging/hst/tracer.json
new file mode 100644
index 0000000..60f4998
--- /dev/null
+++ b/dataset/imaging/hst/tracer.json
@@ -0,0 +1,110 @@
+{
+    "type": "instance",
+    "class_path": "autolens.lens.tracer.Tracer",
+    "arguments": {
+        "galaxies": {
+            "type": "list",
+            "values": [
+                {
+                    "type": "instance",
+                    "class_path": "autogalaxy.galaxy.galaxy.Galaxy",
+                    "arguments": {
+                        "redshift": 0.5,
+                        "label": null,
+                        "bulge": {
+                            "type": "instance",
+                            "class_path": "autogalaxy.profiles.light.standard.sersic.Sersic",
+                            "arguments": {
+                                "intensity": 2.0,
+                                "ell_comps": {
+                                    "type": "tuple",
+                                    "values": [
+                                        0.05263157894736841,
+                                        3.2227547345982974e-18
+                                    ]
+                                },
+                                "centre": {
+                                    "type": "tuple",
+                                    "values": [
+                                        0.0,
+                                        0.0
+                                    ]
+                                },
+                                "sersic_index": 3.0,
+                                "effective_radius": 0.6
+                            }
+                        },
+                        "mass": {
+                            "type": "instance",
+                            "class_path": "autogalaxy.profiles.mass.total.isothermal.Isothermal",
+                            "arguments": {
+                                "ell_comps": {
+                                    "type": "tuple",
+                                    "values": [
+                                        0.05263157894736841,
+                                        3.2227547345982974e-18
+                                    ]
+                                },
+                                "centre": {
+                                    "type": "tuple",
+                                    "values": [
+                                        0.0,
+                                        0.0
+                                    ]
+                                },
+                                "einstein_radius": 1.6
+                            }
+                        },
+                        "shear": {
+                            "type": "instance",
+                            "class_path": "autogalaxy.profiles.mass.sheets.external_shear.ExternalShear",
+                            "arguments": {
+                                "gamma_2": 0.05,
+                                "gamma_1": 0.05
+                            }
+                        }
+                    }
+                },
+                {
+                    "type": "instance",
+                    "class_path": "autogalaxy.galaxy.galaxy.Galaxy",
+                    "arguments": {
+                        "redshift": 1.0,
+                        "label": null,
+                        "bulge": {
+                            "type": "instance",
+                            "class_path": "autogalaxy.profiles.light.standard.sersic_core.SersicCore",
+                            "arguments": {
+                                "intensity": 4.0,
+                                "gamma": 0.25,
+                                "ell_comps": {
+                                    "type": "tuple",
+                                    "values": [
+                                        0.0962250448649376,
+                                        -0.05555555555555551
+                                    ]
+                                },
+                                "centre": {
+                                    "type": "tuple",
+                                    "values": [
+                                        0.0,
+                                        0.0
+                                    ]
+                                },
+                                "sersic_index": 1.0,
+                                "effective_radius": 0.1,
+                                "alpha": 3.0,
+                                "radius_break": 0.025
+                            }
+                        }
+                    }
+                }
+            ]
+        },
+        "cosmology": {
+            "type": "instance",
+            "class_path": "autogalaxy.cosmology.model.Planck15",
+            "arguments": {}
+        }
+    }
+}
\ No newline at end of file
diff --git a/dataset/imaging/source_complex/data.fits b/dataset/imaging/source_complex/data.fits
new file mode 100644
index 0000000..9ae1b79
Binary files /dev/null and b/dataset/imaging/source_complex/data.fits differ
diff --git a/dataset/imaging/source_complex/noise_map.fits b/dataset/imaging/source_complex/noise_map.fits
new file mode 100644
index 0000000..0761473
Binary files /dev/null and b/dataset/imaging/source_complex/noise_map.fits differ
diff --git a/dataset/imaging/source_complex/psf.fits b/dataset/imaging/source_complex/psf.fits
new file mode 100644
index 0000000..b181b8c
Binary files /dev/null and b/dataset/imaging/source_complex/psf.fits differ
diff --git a/dataset/interferometer/sma/data.fits b/dataset/interferometer/sma/data.fits
new file mode 100644
index 0000000..d49ad21
Binary files /dev/null and b/dataset/interferometer/sma/data.fits differ
diff --git a/dataset/interferometer/sma/noise_map.fits b/dataset/interferometer/sma/noise_map.fits
new file mode 100644
index 0000000..f421a36
Binary files /dev/null and b/dataset/interferometer/sma/noise_map.fits differ
diff --git a/dataset/interferometer/sma/positions.json b/dataset/interferometer/sma/positions.json
new file mode 100644
index 0000000..690a98e
--- /dev/null
+++ b/dataset/interferometer/sma/positions.json
@@ -0,0 +1,20 @@
+{
+    "type": "instance",
+    "class_path": "autoarray.structures.grids.irregular_2d.Grid2DIrregular",
+    "arguments": {
+        "values": {
+            "type": "ndarray",
+            "array": [
+                [
+                    -1.35078125,
+                    -0.505632540438727
+                ],
+                [
+                    0.77734375,
+                    1.6648436277960226
+                ]
+            ],
+            "dtype": "float64"
+        }
+    }
+}
\ No newline at end of file
diff --git a/dataset/interferometer/sma/tracer.json b/dataset/interferometer/sma/tracer.json
new file mode 100644
index 0000000..51169ab
--- /dev/null
+++ b/dataset/interferometer/sma/tracer.json
@@ -0,0 +1,111 @@
+{
+    "type": "instance",
+    "class_path": "autolens.lens.tracer.Tracer",
+    "arguments": {
+        "cosmology": {
+            "type": "instance",
+            "class_path": "autogalaxy.cosmology.model.Planck15",
+            "arguments": {}
+        },
+        "galaxies": {
+            "type": "list",
+            "values": [
+                {
+                    "type": "instance",
+                    "class_path": "autogalaxy.galaxy.galaxy.Galaxy",
+                    "arguments": {
+                        "label": null,
+                        "redshift": 0.5,
+                        "mass": {
+                            "type": "instance",
+                            "class_path": "autogalaxy.profiles.mass.total.isothermal.Isothermal",
+                            "arguments": {
+<<<<<<< Updated upstream
+                                "einstein_radius": 1.6,
+                                "centre": {
+=======
+                                "ell_comps": {
+>>>>>>> Stashed changes
+                                    "type": "tuple",
+                                    "values": [
+                                        0.05263157894736841,
+                                        3.2227547345982974e-18
+                                    ]
+                                },
+<<<<<<< Updated upstream
+                                "ell_comps": {
+=======
+                                "centre": {
+>>>>>>> Stashed changes
+                                    "type": "tuple",
+                                    "values": [
+                                        0.0,
+                                        0.0
+                                    ]
+                                },
+                                "einstein_radius": 1.6
+                            }
+                        },
+                        "shear": {
+                            "type": "instance",
+                            "class_path": "autogalaxy.profiles.mass.sheets.external_shear.ExternalShear",
+                            "arguments": {
+                                "gamma_1": 0.05,
+                                "gamma_2": 0.05
+                            }
+                        }
+                    }
+                },
+                {
+                    "type": "instance",
+                    "class_path": "autogalaxy.galaxy.galaxy.Galaxy",
+                    "arguments": {
+                        "label": null,
+                        "redshift": 1.0,
+                        "bulge": {
+                            "type": "instance",
+                            "class_path": "autogalaxy.profiles.light.standard.sersic_core.SersicCore",
+                            "arguments": {
+                                "alpha": 3.0,
+<<<<<<< Updated upstream
+                                "gamma": 0.25,
+                                "effective_radius": 1.0,
+                                "centre": {
+=======
+                                "intensity": 0.3,
+                                "gamma": 0.25,
+                                "radius_break": 0.025,
+                                "ell_comps": {
+>>>>>>> Stashed changes
+                                    "type": "tuple",
+                                    "values": [
+                                        0.1,
+                                        0.1
+                                    ]
+                                },
+<<<<<<< Updated upstream
+                                "intensity": 0.3,
+                                "radius_break": 0.025,
+                                "ell_comps": {
+=======
+                                "centre": {
+>>>>>>> Stashed changes
+                                    "type": "tuple",
+                                    "values": [
+                                        0.0962250448649376,
+                                        -0.05555555555555551
+                                    ]
+                                },
+<<<<<<< Updated upstream
+=======
+                                "effective_radius": 1.0,
+>>>>>>> Stashed changes
+                                "sersic_index": 2.5
+                            }
+                        }
+                    }
+                }
+            ]
+        }
+    }
+}
\ No newline at end of file
diff --git a/dataset/interferometer/sma/uv_wavelengths.fits b/dataset/interferometer/sma/uv_wavelengths.fits
new file mode 100644
index 0000000..02737ba
Binary files /dev/null and b/dataset/interferometer/sma/uv_wavelengths.fits differ
diff --git a/dataset/point_source/simple/point_dataset_positions_only.json b/dataset/point_source/simple/point_dataset_positions_only.json
new file mode 100644
index 0000000..8e02886
--- /dev/null
+++ b/dataset/point_source/simple/point_dataset_positions_only.json
@@ -0,0 +1,56 @@
+{
+    "type": "instance",
+    "class_path": "autolens.point.dataset.PointDataset",
+    "arguments": {
+        "positions_noise_map": {
+            "type": "instance",
+            "class_path": "autoarray.structures.arrays.irregular.ArrayIrregular",
+            "arguments": {
+                "values": {
+                    "type": "ndarray",
+                    "array": [
+                        0.05,
+                        0.05,
+                        0.05,
+                        0.05
+                    ],
+                    "dtype": "float64"
+                }
+            }
+        },
+        "fluxes": null,
+        "time_delays_noise_map": null,
+        "time_delays": null,
+        "positions": {
+            "type": "instance",
+            "class_path": "autoarray.structures.grids.irregular_2d.Grid2DIrregular",
+            "arguments": {
+                "values": {
+                    "type": "ndarray",
+                    "array": [
+                        [
+                            -1.0285552366239088,
+                            -1.092385071774359
+                        ],
+                        [
+                            0.3509304746580634,
+                            1.6278794213275605
+                        ],
+                        [
+                            1.5748094742113987,
+                            0.4342241069464528
+                        ],
+                        [
+                            1.2452712382900812,
+                            1.2238422712179249
+                        ]
+                    ],
+                    "dtype": "float64"
+                }
+            }
+        },
+        "name": "point_0",
+        "fluxes_noise_map": null,
+        "redshift": null
+    }
+}
\ No newline at end of file
diff --git a/dataset/point_source/simple/tracer.json b/dataset/point_source/simple/tracer.json
new file mode 100644
index 0000000..0124d90
--- /dev/null
+++ b/dataset/point_source/simple/tracer.json
@@ -0,0 +1,78 @@
+{
+    "type": "instance",
+    "class_path": "autolens.lens.tracer.Tracer",
+    "arguments": {
+        "galaxies": {
+            "type": "list",
+            "values": [
+                {
+                    "type": "instance",
+                    "class_path": "autogalaxy.galaxy.galaxy.Galaxy",
+                    "arguments": {
+                        "label": null,
+                        "redshift": 0.5,
+                        "mass": {
+                            "type": "instance",
+                            "class_path": "autogalaxy.profiles.mass.total.isothermal.Isothermal",
+                            "arguments": {
+                                "centre": {
+                                    "type": "tuple",
+                                    "values": [
+                                        0.0,
+                                        0.0
+                                    ]
+                                },
+                                "ell_comps": {
+                                    "type": "tuple",
+                                    "values": [
+                                        0.05263157894736841,
+                                        3.2227547345982974e-18
+                                    ]
+                                },
+                                "einstein_radius": 1.6
+                            }
+                        }
+                    }
+                },
+                {
+                    "type": "instance",
+                    "class_path": "autogalaxy.galaxy.galaxy.Galaxy",
+                    "arguments": {
+                        "label": null,
+                        "redshift": 1.0,
+                        "light": {
+                            "type": "instance",
+                            "class_path": "autogalaxy.profiles.light.standard.exponential_core.ExponentialCore",
+                            "arguments": {
+                                "alpha": 3.0,
+                                "intensity": 0.1,
+                                "centre": {
+                                    "type": "tuple",
+                                    "values": [
+                                        0.07,
+                                        0.07
+                                    ]
+                                },
+                                "radius_break": 0.025,
+                                "gamma": 0.25,
+                                "ell_comps": {
+                                    "type": "tuple",
+                                    "values": [
+                                        0.0,
+                                        0.0
+                                    ]
+                                },
+                                "effective_radius": 0.02
+                            }
+                        }
+                    }
+                }
+            ]
+        },
+        "cosmology": {
+            "type": "instance",
+            "class_path": "autogalaxy.cosmology.model.Planck15",
+            "arguments": {}
+        }
+    }
+}
\ No newline at end of file
diff --git a/likelihood/README.md b/likelihood/README.md
index 25bd058..0751281 100644
--- a/likelihood/README.md
+++ b/likelihood/README.md
@@ -1,7 +1,63 @@
 # likelihood
 
-Likelihood JIT profiling scripts and per-section READMEs.
+JAX JIT profiling for the PyAutoLens likelihood function across imaging, interferometer, point-source, and datacube datasets, and across the MGE, pixelization, and Delaunay model compositions used in real science cases.
 
-Populated by **Phase 1** of the `autolens_profiling` roadmap. Will mirror `autolens_workspace_developer/jax_profiling/jit/{imaging,interferometer,point_source,datacube}/`, preserving the step-by-step profile narrative each script already prints.
+## What "JIT likelihood profiling" means
 
-See the top-level [README](../README.md) for the full phase plan.
+For each science case, the likelihood function turns a parameter vector into a single number (log-likelihood) via a chain of array operations: instantiate the model, build the `Tracer`, ray-trace grids through the lens, compute a mapping matrix, blur it with the PSF, solve a linear-algebra reconstruction problem, and finally compute a chi-squared. Under `xp=jnp`, every step is dispatched as a JAX op and can be compiled into a single XLA program with `jax.jit`.
+
+Profiling the **whole likelihood** as one JIT'd function gives the honest per-call cost a sampler will see in production. Profiling **each step individually** under its own JIT gives the breakdown that tells you where the time is going. Both numbers matter: the whole-function timing is the production cost, and the per-step breakdown is the optimisation target. Each script in this section reports both wherever the underlying pipeline supports per-step JIT-ing (the interferometer and datacube paths intentionally stay at full-pipeline JIT for now — see those subfolders' READMEs for why).
+
+Every script also reports a **batched (`jax.vmap`) per-likelihood cost** to make explicit how much the JIT amortises across a population of evaluations — the regime an actual sampler operates in.
+
+## How to read the per-script output
+
+Each script prints a structured narrative to stdout, ending in:
+
+- The eager (numpy) baseline log-likelihood for sanity.
+- The single-JIT lower / compile / first-call / steady-state per-call timings.
+- The vmap per-likelihood cost and speedup vs single-JIT.
+- A correctness check: eager ≡ JIT ≡ vmap log-likelihoods at `rtol=1e-4`.
+- A `results/likelihood/<type>/<script>_likelihood_summary_<instrument>_v<al.__version__>.{json,png}` write.
+
+The JSON carries the structured timings keyed by step name plus the model / dataset metadata. The PNG is a bar chart of per-step costs (where applicable) plus the single-JIT vs vmap comparison.
+
+## Versioned artifacts
+
+Result files are tagged with the PyAutoLens release that produced them (`al.__version__`). Old versions remain alongside new ones so cross-release trends stay visible — Phase 4's dashboard will read the latest per axis and present the headline numbers framed by astronomy instrument.
+
+See the top-level [results/README.md](../results/README.md) for the full filename convention.
+
+## Sections
+
+| Folder | Profiles |
+|--------|----------|
+| [`imaging/`](./imaging/README.md) | MGE, pixelization, and Delaunay likelihoods on imaging datasets (HST-resolution by default). |
+| [`interferometer/`](./interferometer/README.md) | MGE, pixelization, and Delaunay likelihoods on interferometer (visibility-space) datasets (SMA by default). |
+| [`point_source/`](./point_source/README.md) | Image-plane and source-plane chi-squared for lensed point sources. |
+| [`datacube/`](./datacube/README.md) | Multi-channel datacube likelihoods (e.g. ALMA-style) with Delaunay pixelization. |
+
+## Running a script
+
+From the repo root:
+
+```bash
+cd autolens_profiling
+python likelihood/imaging/mge.py
+```
+
+Scripts use the input datasets under `dataset/<type>/<instrument>/` (see top-level [README](../README.md)). The default instrument is encoded per-script; some support a CLI flag to switch instruments. Run with `--help` for the supported options.
+
+**Codex / sandboxed runs** — set writable cache dirs so numba and matplotlib don't choke on read-only home/source paths:
+
+```bash
+NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib python likelihood/imaging/mge.py
+```
+
+## Conventions inherited from `autolens_workspace_developer`
+
+The scripts follow the JIT conventions documented at `autolens_workspace_developer/CLAUDE.md`:
+
+- All autoarray types (`Array2D`, `Grid2D`, `Grid2DIrregular`, …) expose `.array` for the raw `np.ndarray` / `jax.Array` underneath. These are extracted before crossing the `jax.jit` boundary because autoarray types are not registered as JAX pytrees as **inputs**.
+- The `xp` parameter (`xp=np` default, `xp=jnp` for JAX) controls the backend. JIT'd closures pass `xp=jnp` through every nested call.
+- The model is converted to a JAX pytree via `autofit.jax.register_model(model)` so `af.ModelInstance` can cross the JIT boundary directly — no manual flat-vector unpacking.
diff --git a/likelihood/datacube/README.md b/likelihood/datacube/README.md
new file mode 100644
index 0000000..a05c16a
--- /dev/null
+++ b/likelihood/datacube/README.md
@@ -0,0 +1,55 @@
+# likelihood/datacube
+
+JAX JIT profiling for the PyAutoLens **datacube** likelihood function — fitting an N-channel cube of interferometer observations (e.g. ALMA velocity / frequency channels) that share a single lens model, with each channel reconstructing its own Delaunay-pixelized source.
+
+## Channel-invariant vs channel-variant split
+
+The key new ingredient relative to the single-channel interferometer path is the **channel-invariant vs channel-variant** decomposition: most steps are computed once for the whole cube (shared lens model, shared mesh, shared mask), while only a few steps recur per channel.
+
+| Step | Channel-invariant? | Computed |
+|------|--------------------|----------|
+| 1. Ray-trace data grid | yes | once for the cube |
+| 2. Ray-trace mesh grid | yes | once for the cube |
+| 3. Inversion setup (border + mapper + NUFFT) | NUFFT depends on `uv_wavelengths` | once per channel |
+| 4. Data vector D | per channel | once per channel |
+| 5. Curvature matrix F | per channel | once per channel |
+| 6. Regularization matrix H | yes | once for the cube |
+| 7. Reconstruction (NNLS) | per channel | once per channel |
+| 8. Mapped recon + log-evidence | per channel | once per channel |
+
+The cube total is:
+
+```
+cube_cost = sum(channel_invariant_costs) + N_channels * sum(channel_variant_costs)
+```
+
+That number quantifies how much a future "shared `Lᵀ W̃ L`" optimisation would save: moving the curvature matrix from per-channel to shared would subtract `(N − 1) * curvature_matrix_cost` from the cube total. The profiling script reports this number directly.
+
+## Scripts
+
+| Script | What it profiles |
+|--------|------------------|
+| [`delaunay.py`](./delaunay.py) | Step-by-step JIT profiling of an N-channel datacube with shared lens model and per-channel Delaunay source reconstruction. Mirrors the per-step structure of [`../interferometer/delaunay.py`](../interferometer/delaunay.py). |
+
+## Default dataset
+
+`dataset/interferometer/sma/` — the same SMA-like mock used by the single-channel interferometer scripts, **loaded N times as a 4-channel cube**. Every channel has identical visibilities, noise map, and uv_wavelengths — the point here is timing, not science. The N-channel cube log-evidence is exactly `N × single-channel log-evidence`, which makes the regression assertion trivial.
+
+For a realistic per-channel-distinct cube, point the loader at the workspace simulator output at `autolens_workspace/dataset/interferometer/datacube/sim_simple/`. The JIT-cost taxonomy doesn't change — it's a function of which arrays are loop-variables in `FitInterferometer`, not the data values themselves.
+
+## Headline run-times (populated by Phase 4)
+
+| Script | Dataset | N channels | CPU | Laptop GPU | A100 |
+|--------|---------|------------|-----|------------|------|
+| `delaunay.py` | SMA × 4 | 4 | _populated_ | _populated_ | _populated_ |
+
+Numbers are the **steady-state per-call cost** (single-JIT, post-warmup), in milliseconds. Phase 4's dashboard auto-fills this from the latest `*_summary_v<version>.json` artifacts under `results/likelihood/datacube/`.
+
+## Output
+
+The script writes:
+
+```
+results/likelihood/datacube/delaunay_likelihood_summary_<instrument>_v<al.__version__>.json
+results/likelihood/datacube/delaunay_likelihood_summary_<instrument>_v<al.__version__>.png
+```
diff --git a/likelihood/datacube/__init__.py b/likelihood/datacube/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/likelihood/datacube/delaunay.py b/likelihood/datacube/delaunay.py
new file mode 100644
index 0000000..67c3a81
--- /dev/null
+++ b/likelihood/datacube/delaunay.py
@@ -0,0 +1,981 @@
+"""
+JAX Profiling: Delaunay Datacube Likelihood (Step-by-Step)
+==========================================================
+
+Profiles each step of the JAX likelihood function for an ALMA-style datacube —
+a list of N ``Interferometer`` channels sharing a single lens model — where
+each channel reconstructs its own source with a Delaunay pixelization +
+ConstantSplit regularization.
+
+Mirrors the step-by-step structure of
+``likelihood/interferometer/delaunay.py`` (Phase 2 of the datacube
+roadmap, just merged). The key new ingredient is the **channel-invariant vs
+channel-variant** split: most steps are computed once for the whole cube
+(shared lens, shared mesh, shared mask), only the NUFFT-based inversion-setup
+chain, the data vector, the curvature matrix, the reconstruction, and the
+log-evidence depend on per-channel data.
+
+The cube total is::
+
+    cube_cost = sum(channel_invariant_costs) + N_channels * sum(channel_variant_costs)
+
+That number quantifies how much the deferred shared-``Lᵀ W̃ L`` optimisation
+will save: moving the curvature matrix from per-channel to shared would
+subtract ``(N - 1) * curvature_matrix_cost`` from the cube total.
+
+Channel-invariant vs channel-variant taxonomy
+---------------------------------------------
+
+For the canonical datacube case where the lens model is shared across all
+channels:
+
+============================================  ================  =========================
+Step                                          Channel-invariant Computed
+============================================  ================  =========================
+1. Ray-trace data grid                        yes               once for the cube
+2. Ray-trace mesh grid                        yes               once for the cube
+3. Inversion setup (border + mapper + NUFFT)  **NUFFT depends   once per channel
+                                              on uv_wavelengths**
+4. Data vector D                              per channel       once per channel
+5. Curvature matrix F                         per channel       once per channel
+6. Regularization matrix H                    yes               once for the cube
+7. Reconstruction (NNLS)                      per channel       once per channel
+8. Mapped recon + log evidence                per channel       once per channel
+============================================  ================  =========================
+
+Dataset
+-------
+
+This profiler reuses the SMA interferometer dataset
+(``dataset/interferometer/sma/``) loaded N times as a 4-channel
+"cube". Each channel has identical visibilities, noise map and uv_wavelengths
+— the point here is timing, not science. The N-channel cube log-evidence is
+``N × single-channel log-evidence`` exactly, which makes the regression
+assertion trivial.
+
+If you want a realistic per-channel-distinct cube, point the loader at the
+workspace simulator output at
+``../autolens_workspace/dataset/interferometer/datacube/sim_simple/``; the
+JIT-cost taxonomy doesn't change because it's a function of which arrays are
+loop-variables in ``FitInterferometer``, not the data values themselves.
+
+Measures
+--------
+
+1. Eager baseline: ``FitInterferometer`` per channel with ``xp=np``; cube
+   reference log-evidence is the sum.
+2. Per-step JIT profiling: each pipeline stage gets its own ``jit_profile()``
+   call (lower / compile / first-call / steady-state × 10). Channel-invariant
+   stages are timed once; channel-variant stages are timed on channel 0 and
+   the cube cost is reported as ``N × per-call``.
+3. Full-pipeline cube JIT: ``jax.jit`` over the explicit
+   ``sum(analysis.log_likelihood_function(instance) for analysis in
+   analysis_list)`` — the same shape as the user-facing
+   ``datacube/likelihood_function.py`` and the cube modeling scripts'
+   internal ``FactorGraphModel`` sum.
+4. Correctness: per-step recomputed cube log-evidence and full-pipeline JIT
+   log-evidence both match the summed eager ``FitInterferometer.log_evidence``
+   at ``rtol=1e-4``.
+5. Results JSON + bar chart written to ``results/jit/datacube/`` using the
+   same schema as the interferometer sibling. Bar chart shows the cube-total
+   form of every step (channel-variant entries pre-multiplied by N).
+
+vmap is **skipped** for the cube profiler. The natural batching dimension is
+"datasets" (one entry per channel) not "parameters" (which the
+interferometer-sibling vmap exercises). A vmap-over-channels variant would
+require a different graph shape and isn't the bottleneck we care about for
+the shared-``Lᵀ W̃ L`` optimisation.
+"""
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import time
+import subprocess
+import sys
+from pathlib import Path
+from contextlib import contextmanager
+
+import autofit as af
+import autolens as al
+import autoarray as aa
+from autofit.jax import register_model as _register_model_pytrees
+
+# ---------------------------------------------------------------------------
+# Instrument configuration
+# ---------------------------------------------------------------------------
+
+INSTRUMENTS = {
+    "sma": {"pixel_scale": 0.1, "real_space_shape": (256, 256)},
+    "alma": {"pixel_scale": 0.05, "real_space_shape": (256, 256)},
+}
+
+instrument = "sma"  # <-- change this to profile a different instrument
+
+n_channels = 4
+overlay_shape = (26, 26)
+edge_n_points = 30
+regularization_coefficient = 1.0
+
+
+# ---------------------------------------------------------------------------
+# Profiling helpers
+# ---------------------------------------------------------------------------
+
+class Timer:
+    """Accumulates named timing measurements and prints a summary."""
+
+    def __init__(self):
+        self.records: list[tuple[str, float]] = []
+
+    @contextmanager
+    def section(self, label: str):
+        start = time.perf_counter()
+        yield
+        elapsed = time.perf_counter() - start
+        self.records.append((label, elapsed))
+        print(f"  [{label}] {elapsed:.4f} s")
+
+
+def block(x):
+    if hasattr(x, "block_until_ready"):
+        x.block_until_ready()
+    return x
+
+
+def jit_profile(func, label, *args, n_repeats=10):
+    """JIT-compile *func*, time lower / compile / first call / steady state."""
+    jitted = jax.jit(func)
+
+    with timer.section(f"{label}_lower"):
+        lowered = jitted.lower(*args)
+
+    with timer.section(f"{label}_compile"):
+        compiled = lowered.compile()
+
+    with timer.section(f"{label}_first_call"):
+        result = compiled(*args)
+        block(result)
+
+    with timer.section(f"{label}_steady_x{n_repeats}"):
+        for _ in range(n_repeats):
+            result = compiled(*args)
+            block(result)
+
+    per_call = timer.records[-1][1] / n_repeats
+    print(f"    -> per-call avg: {per_call:.6f} s")
+    return compiled, result
+
+
+timer = Timer()
+likelihood_steps = []  # (label, per_call_seconds) — cube-total cost for each step
+
+# ===================================================================
+# PART A — Setup (not JIT-compiled)
+# ===================================================================
+
+# ---------------------------------------------------------------------------
+# 1. Dataset loading: reuse SMA interferometer dataset N times
+# ---------------------------------------------------------------------------
+
+print(f"\n--- Dataset loading [{instrument}, {n_channels} channels] ---")
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[1]
+pixel_scale = INSTRUMENTS[instrument]["pixel_scale"]
+real_space_shape = INSTRUMENTS[instrument]["real_space_shape"]
+dataset_path = Path("dataset") / "interferometer" / instrument
+
+if al.util.dataset.should_simulate(str(dataset_path)):
+    raise FileNotFoundError(
+        f"Input dataset missing at '{dataset_path}'. The autolens_profiling "
+        f"repo mirrors only the curated datasets needed for default smoke "
+        f"runs. To regenerate or extend datasets, use the source-of-truth "
+        f"scripts under autolens_workspace_developer/jax_profiling/dataset_setup/, "
+        f"then copy the result into autolens_profiling/dataset/."
+    )
+
+mask_radius = 3.0
+
+real_space_mask = al.Mask2D.circular(
+    shape_native=real_space_shape,
+    pixel_scales=pixel_scale,
+    radius=mask_radius,
+)
+
+with timer.section("dataset_list_load"):
+    dataset_list = [
+        al.Interferometer.from_fits(
+            data_path=dataset_path / "data.fits",
+            noise_map_path=dataset_path / "noise_map.fits",
+            uv_wavelengths_path=dataset_path / "uv_wavelengths.fits",
+            real_space_mask=real_space_mask,
+            transformer_class=al.TransformerDFT,
+        )
+        for _ in range(n_channels)
+    ]
+
+n_visibilities = dataset_list[0].uv_wavelengths.shape[0]
+print(f"  Channels:           {n_channels}")
+print(f"  Visibilities/chan:  {n_visibilities}")
+
+# ---------------------------------------------------------------------------
+# 2. Image mesh + edge points (Delaunay-specific, channel-invariant)
+# ---------------------------------------------------------------------------
+
+print("\n--- Image mesh construction (Delaunay) ---")
+
+with timer.section("image_mesh_overlay"):
+    image_mesh = al.image_mesh.Overlay(shape=overlay_shape)
+    image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(
+        mask=dataset_list[0].real_space_mask
+    )
+
+with timer.section("edge_points"):
+    pre_edge_pixels = image_plane_mesh_grid.shape[0]
+    image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(
+        image_plane_mesh_grid=image_plane_mesh_grid,
+        centre=(0.0, 0.0),
+        radius=mask_radius,
+        n_points=edge_n_points,
+    )
+    edge_pixels_total = image_plane_mesh_grid.shape[0] - pre_edge_pixels
+
+n_mesh_vertices = image_plane_mesh_grid.shape[0]
+print(f"  Overlay shape: {overlay_shape}")
+print(f"  Mesh vertices (incl. edge): {n_mesh_vertices}")
+print(f"  Edge points added: {edge_pixels_total}")
+
+# ---------------------------------------------------------------------------
+# 3. Model construction
+# ---------------------------------------------------------------------------
+
+print("\n--- Model construction ---")
+
+with timer.section("model_build"):
+    mass = af.Model(al.mp.Isothermal)
+    mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.005)
+    mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.005)
+    mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05)
+    _lens_mass_ell = al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0)
+    mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=_lens_mass_ell[0], sigma=0.01)
+    mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=_lens_mass_ell[1], sigma=0.01)
+
+    shear = af.Model(al.mp.ExternalShear)
+    shear.gamma_1 = af.GaussianPrior(mean=0.05, sigma=0.005)
+    shear.gamma_2 = af.GaussianPrior(mean=0.05, sigma=0.005)
+
+    lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)
+
+    mesh = al.mesh.Delaunay(
+        pixels=n_mesh_vertices,
+        zeroed_pixels=edge_pixels_total,
+    )
+    regularization = al.reg.ConstantSplit(coefficient=regularization_coefficient)
+    pixelization = al.Pixelization(mesh=mesh, regularization=regularization)
+
+    source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)
+
+    model = af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+print(f"  Total free parameters: {model.total_free_parameters}")
+print(f"  Delaunay pixels: {n_mesh_vertices}")
+
+# ---------------------------------------------------------------------------
+# 4. Instantiate concrete objects from prior medians
+# ---------------------------------------------------------------------------
+
+print("\n--- Instantiate concrete model ---")
+
+with timer.section("instance_from_vector"):
+    param_vector = model.physical_values_from_prior_medians
+    instance = model.instance_from_vector(vector=param_vector)
+
+with timer.section("register_pytrees"):
+    _register_model_pytrees(model)
+
+params_tree = jax.tree_util.tree_map(jnp.asarray, instance)
+
+tracer = al.Tracer(galaxies=list(instance.galaxies))
+
+# The adapt_images object is channel-invariant — the image-plane Delaunay mesh
+# vertices are shared across channels (the lens model is shared).
+adapt_images = al.AdaptImages(
+    galaxy_image_plane_mesh_grid_dict={
+        instance.galaxies.source: image_plane_mesh_grid,
+    },
+    galaxy_name_image_plane_mesh_grid_dict={
+        "('galaxies', 'source')": image_plane_mesh_grid,
+    },
+)
+
+print(f"  Tracer planes: {tracer.total_planes}")
+
+# ---------------------------------------------------------------------------
+# 5. Configuration summary
+# ---------------------------------------------------------------------------
+
+print("\n--- Configuration (determines run time) ---")
+print(f"  Instrument:              {instrument}")
+print(f"  Channels:                {n_channels}")
+print(f"  Pixel scale:             {pixel_scale} arcsec/pixel")
+print(f"  Real-space mask radius:  {mask_radius} arcsec")
+print(f"  Real-space grid shape:   {real_space_shape[0]} x {real_space_shape[1]}")
+print(f"  Visibilities/chan:       {n_visibilities}")
+print(f"  Overlay shape:           {overlay_shape[0]} x {overlay_shape[1]}")
+print(f"  Delaunay vertices:       {n_mesh_vertices}")
+print(f"  Edge zeroed pixels:      {edge_pixels_total}")
+print(f"  Reg. coefficient:        {regularization_coefficient}")
+
+# ---------------------------------------------------------------------------
+# 6. Per-channel eager FitInterferometer baseline
+# ---------------------------------------------------------------------------
+
+print(f"\n--- Per-channel eager FitInterferometer baselines ({n_channels} channels) ---")
+
+fit_list = []
+log_evidence_per_channel = []
+with timer.section(f"eager_fit_per_channel_x{n_channels}"):
+    for c, dataset in enumerate(dataset_list):
+        f = al.FitInterferometer(
+            dataset=dataset,
+            tracer=tracer,
+            adapt_images=adapt_images,
+            xp=np,
+        )
+        fit_list.append(f)
+        log_evidence_per_channel.append(f.log_evidence)
+
+for c, le in enumerate(log_evidence_per_channel):
+    print(f"  channel {c}: log_evidence = {le:.6f}")
+
+cube_log_evidence_ref = float(sum(log_evidence_per_channel))
+print(f"  cube reference log_evidence (sum) = {cube_log_evidence_ref:.6f}")
+
+
+# ===================================================================
+# PART B — Per-step JIT profiling (channel-invariant + channel-variant)
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("PER-STEP JIT PROFILING — CUBE")
+print("=" * 70)
+print(f"  Channel-invariant steps are timed once.")
+print(f"  Channel-variant steps are JIT-compiled on channel 0; the reported")
+print(f"  cube cost is N × the per-channel steady-state per-call.")
+
+# Reference single-channel context (channel 0)
+fit = fit_list[0]
+dataset = dataset_list[0]
+
+# Extract raw arrays from autoarray types via .array so they can cross
+# JIT boundaries. See CLAUDE.md for rationale.
+grid_pix_raw = jnp.array(dataset.grids.pixelization.array)
+mesh_grid_raw = jnp.array(image_plane_mesh_grid.array)
+
+# ---------------------------------------------------------------------------
+# Step 1: Ray-trace data grid (channel-invariant)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 1: Ray-trace data grid (shared) ---")
+
+with timer.section("ray_trace_data_eager"):
+    traced_grids = tracer.traced_grid_2d_list_from(
+        grid=dataset.grids.pixelization, xp=jnp
+    )
+    for tg in traced_grids:
+        block(tg)
+
+
+def ray_trace_data_raw(grid_raw):
+    grid = aa.Grid2DIrregular(values=grid_raw, xp=jnp)
+    traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp)
+    return jnp.stack([tg.array for tg in traced])
+
+
+_, _ = jit_profile(ray_trace_data_raw, "ray_trace_data_jit", grid_pix_raw)
+ray_trace_data_per_call = timer.records[-1][1] / 10
+likelihood_steps.append(
+    ("Ray-trace data grid (shared)", ray_trace_data_per_call)
+)
+
+# ---------------------------------------------------------------------------
+# Step 2: Ray-trace mesh grid (channel-invariant)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 2: Ray-trace mesh grid (shared) ---")
+
+
+def ray_trace_mesh_raw(mesh_raw):
+    grid = aa.Grid2DIrregular(values=mesh_raw, xp=jnp)
+    traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp)
+    return jnp.stack([tg.array for tg in traced])
+
+
+_, _ = jit_profile(ray_trace_mesh_raw, "ray_trace_mesh_jit", mesh_grid_raw)
+ray_trace_mesh_per_call = timer.records[-1][1] / 10
+likelihood_steps.append(
+    ("Ray-trace mesh grid (shared)", ray_trace_mesh_per_call)
+)
+
+# ---------------------------------------------------------------------------
+# Extract inversion matrices from channel 0
+# ---------------------------------------------------------------------------
+
+print("\n--- Extracting inversion matrices from channel 0 ---")
+
+inversion = fit.inversion
+
+with timer.section("extract_inversion_matrices"):
+    transformed_mm_ref = jnp.asarray(inversion.operated_mapping_matrix)
+    mapping_matrix_ref = jnp.asarray(inversion.mapping_matrix)
+
+    inv_mapper = inversion.cls_list_from(cls=al.Mapper)[0]
+    neighbors = inv_mapper.neighbors
+    neighbors_array = jnp.array(np.asarray(neighbors))
+    neighbors_sizes = jnp.array(neighbors.sizes)
+
+print(f"  transformed_mapping_matrix shape: {transformed_mm_ref.shape}")
+print(f"  transformed_mapping_matrix dtype: {transformed_mm_ref.dtype}")
+print(f"  mapping_matrix shape: {mapping_matrix_ref.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 3: Inversion setup (per channel — NUFFT depends on uv_wavelengths)
+# ---------------------------------------------------------------------------
+# Steps 5-8 from the interferometer-sibling numbering (border + Delaunay +
+# mapper + mapping matrix + NUFFT), combined and JIT-profiled from a pytree
+# ModelInstance. Channel-variant because each channel's NUFFT uses its own
+# uv_wavelengths. JIT-compile on channel 0; report cube cost as N × per-call.
+
+print("\n--- Step 3: Inversion setup, incl. NUFFT (per channel) ---")
+
+
+def transformed_mm_from_params(params_tree):
+    """Inversion setup from a pytree ModelInstance — full chain through NUFFT.
+
+    This closes over ``dataset`` (channel 0) for the JIT compilation. In real
+    cube usage each channel's `AnalysisFactor` closes over its own
+    `dataset`, so the steady-state per-call cost is what we want to scale by N.
+    """
+    t = al.Tracer(galaxies=list(params_tree.galaxies))
+    adapt_images_jax = al.AdaptImages(
+        galaxy_image_plane_mesh_grid_dict={
+            params_tree.galaxies.source: image_plane_mesh_grid,
+        },
+        galaxy_name_image_plane_mesh_grid_dict={
+            "('galaxies', 'source')": image_plane_mesh_grid,
+        },
+    )
+    fit_jax = al.FitInterferometer(
+        dataset=dataset,
+        tracer=t,
+        adapt_images=adapt_images_jax,
+        xp=jnp,
+    )
+    return jnp.asarray(fit_jax.inversion.operated_mapping_matrix)
+
+
+_, transformed_mm_jit = jit_profile(
+    transformed_mm_from_params, "inversion_setup_jit", params_tree
+)
+inversion_setup_per_channel = timer.records[-1][1] / 10
+likelihood_steps.append(
+    (
+        f"Inversion setup, incl. NUFFT (per channel × {n_channels})",
+        n_channels * inversion_setup_per_channel,
+    )
+)
+
+print(f"  per-channel: {inversion_setup_per_channel:.6f} s")
+print(f"  cube cost (× {n_channels}): {n_channels * inversion_setup_per_channel:.6f} s")
+
+# Use the reference real / imag arrays for the linear-algebra steps
+transformed_mm_real_jnp = jnp.real(transformed_mm_ref)
+transformed_mm_imag_jnp = jnp.imag(transformed_mm_ref)
+data_real_jnp = jnp.array(dataset.data.real)
+data_imag_jnp = jnp.array(dataset.data.imag)
+noise_real_jnp = jnp.array(dataset.noise_map.real)
+noise_imag_jnp = jnp.array(dataset.noise_map.imag)
+
+# ---------------------------------------------------------------------------
+# Step 4: Data vector D (per channel)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 4: Data vector D (per channel) ---")
+
+
+def compute_data_vector(
+    transformed_mm_real, transformed_mm_imag, data_real, data_imag,
+    noise_real, noise_imag,
+):
+    weighted_data_real = data_real / (noise_real ** 2)
+    weighted_data_imag = data_imag / (noise_imag ** 2)
+    return jnp.matmul(transformed_mm_real.T, weighted_data_real) + jnp.matmul(
+        transformed_mm_imag.T, weighted_data_imag
+    )
+
+
+with timer.section("data_vector_eager"):
+    data_vector = compute_data_vector(
+        transformed_mm_real_jnp, transformed_mm_imag_jnp,
+        data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp,
+    )
+    block(data_vector)
+
+_, data_vector = jit_profile(
+    compute_data_vector, "data_vector_jit",
+    transformed_mm_real_jnp, transformed_mm_imag_jnp,
+    data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp,
+)
+data_vector_per_channel = timer.records[-1][1] / 10
+likelihood_steps.append(
+    (
+        f"Data vector D (per channel × {n_channels})",
+        n_channels * data_vector_per_channel,
+    )
+)
+
+# ---------------------------------------------------------------------------
+# Step 5: Curvature matrix F (per channel)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 5: Curvature matrix F (per channel) ---")
+
+no_reg_list = list(inversion.no_regularization_index_list)
+
+
+def compute_curvature_matrix(
+    transformed_mm_real, transformed_mm_imag, noise_real, noise_imag,
+):
+    real_curv = al.util.inversion.curvature_matrix_via_mapping_matrix_from(
+        mapping_matrix=transformed_mm_real,
+        noise_map=noise_real,
+        settings=fit.settings,
+        add_to_curvature_diag=True,
+        no_regularization_index_list=no_reg_list,
+        xp=jnp,
+    )
+    imag_curv = al.util.inversion.curvature_matrix_via_mapping_matrix_from(
+        mapping_matrix=transformed_mm_imag,
+        noise_map=noise_imag,
+        settings=fit.settings,
+        add_to_curvature_diag=False,
+        no_regularization_index_list=no_reg_list,
+        xp=jnp,
+    )
+    return real_curv + imag_curv
+
+
+with timer.section("curvature_matrix_eager"):
+    curvature_matrix = compute_curvature_matrix(
+        transformed_mm_real_jnp, transformed_mm_imag_jnp, noise_real_jnp, noise_imag_jnp,
+    )
+    block(curvature_matrix)
+
+_, curvature_matrix = jit_profile(
+    compute_curvature_matrix, "curvature_matrix_jit",
+    transformed_mm_real_jnp, transformed_mm_imag_jnp, noise_real_jnp, noise_imag_jnp,
+)
+curvature_matrix_per_channel = timer.records[-1][1] / 10
+likelihood_steps.append(
+    (
+        f"Curvature matrix F (per channel × {n_channels})",
+        n_channels * curvature_matrix_per_channel,
+    )
+)
+
+# ---------------------------------------------------------------------------
+# Step 6: Regularization matrix H (channel-invariant)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 6: Regularization matrix H (shared) ---")
+
+with timer.section("regularization_matrix_eager"):
+    regularization_matrix = jnp.array(inversion.regularization_matrix)
+    block(regularization_matrix)
+
+likelihood_steps.append(
+    ("Regularization matrix H (shared)", timer.records[-1][1])
+)
+
+# ---------------------------------------------------------------------------
+# Step 7: Reconstruction NNLS (per channel)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 7: Regularized reconstruction (per channel) ---")
+
+
+def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix):
+    curvature_reg_matrix = curvature_matrix + regularization_matrix
+    return al.util.inversion.reconstruction_positive_only_from(
+        data_vector=data_vector,
+        curvature_reg_matrix=curvature_reg_matrix,
+        xp=jnp,
+    )
+
+
+with timer.section("reconstruction_eager"):
+    reconstruction = compute_reconstruction(
+        jnp.array(data_vector),
+        jnp.array(curvature_matrix),
+        jnp.array(regularization_matrix),
+    )
+    block(reconstruction)
+
+_, reconstruction = jit_profile(
+    compute_reconstruction, "reconstruction_jit",
+    jnp.array(data_vector), jnp.array(curvature_matrix), jnp.array(regularization_matrix),
+)
+reconstruction_per_channel = timer.records[-1][1] / 10
+likelihood_steps.append(
+    (
+        f"Reconstruction NNLS (per channel × {n_channels})",
+        n_channels * reconstruction_per_channel,
+    )
+)
+
+# ---------------------------------------------------------------------------
+# Step 8: Mapped recon + log evidence (per channel)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 8: Mapped recon + log evidence (per channel) ---")
+
+
+def compute_log_evidence(
+    data_real, data_imag, noise_real, noise_imag,
+    transformed_mm_real, transformed_mm_imag,
+    reconstruction, curvature_matrix, regularization_matrix, mapper_indices,
+):
+    mapped_real = jnp.matmul(transformed_mm_real, reconstruction)
+    mapped_imag = jnp.matmul(transformed_mm_imag, reconstruction)
+
+    chi_real = jnp.sum(((data_real - mapped_real) / noise_real) ** 2)
+    chi_imag = jnp.sum(((data_imag - mapped_imag) / noise_imag) ** 2)
+    chi_squared = chi_real + chi_imag
+
+    regularization_term = jnp.dot(
+        reconstruction, jnp.dot(regularization_matrix, reconstruction)
+    )
+
+    curvature_reg_matrix = curvature_matrix + regularization_matrix
+    creg_reduced = curvature_reg_matrix[mapper_indices][:, mapper_indices]
+    reg_reduced = regularization_matrix[mapper_indices][:, mapper_indices]
+    log_det_curvature_reg = 2.0 * jnp.sum(
+        jnp.log(jnp.diag(jnp.linalg.cholesky(creg_reduced)))
+    )
+    log_det_regularization = 2.0 * jnp.sum(
+        jnp.log(jnp.diag(jnp.linalg.cholesky(reg_reduced)))
+    )
+
+    noise_normalization = (
+        jnp.sum(jnp.log(2 * jnp.pi * noise_real ** 2))
+        + jnp.sum(jnp.log(2 * jnp.pi * noise_imag ** 2))
+    )
+
+    return -0.5 * (
+        chi_squared + regularization_term + log_det_curvature_reg
+        - log_det_regularization + noise_normalization
+    )
+
+
+mapper_indices_jnp = jnp.array(np.asarray(inversion.mapper_indices))
+inv_recon_jnp = jnp.asarray(inversion.reconstruction)
+inv_curv_jnp = jnp.asarray(inversion.curvature_matrix)
+reg_jnp = jnp.array(regularization_matrix)
+
+with timer.section("log_evidence_eager"):
+    log_evidence_one_channel = compute_log_evidence(
+        data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp,
+        transformed_mm_real_jnp, transformed_mm_imag_jnp,
+        reconstruction, curvature_matrix, reg_jnp, mapper_indices_jnp,
+    )
+    block(log_evidence_one_channel)
+
+_, log_evidence_one_channel = jit_profile(
+    compute_log_evidence, "log_evidence_jit",
+    data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp,
+    transformed_mm_real_jnp, transformed_mm_imag_jnp,
+    reconstruction, curvature_matrix, reg_jnp, mapper_indices_jnp,
+)
+log_evidence_per_channel_cost = timer.records[-1][1] / 10
+likelihood_steps.append(
+    (
+        f"Mapped recon + log evidence (per channel × {n_channels})",
+        n_channels * log_evidence_per_channel_cost,
+    )
+)
+
+print(f"  channel 0 log_evidence (step-by-step) = {log_evidence_one_channel}")
+
+# Correctness check: recompute per-channel log_evidence using each channel's
+# inversion matrices and sum to get the cube log-evidence. Should match the
+# summed eager FitInterferometer.log_evidence at rtol=1e-4.
+log_evidence_check_per_channel = []
+for c, f in enumerate(fit_list):
+    inv_c = f.inversion
+    tm_real = jnp.real(jnp.asarray(inv_c.operated_mapping_matrix))
+    tm_imag = jnp.imag(jnp.asarray(inv_c.operated_mapping_matrix))
+    le_c = compute_log_evidence(
+        jnp.array(dataset_list[c].data.real),
+        jnp.array(dataset_list[c].data.imag),
+        jnp.array(dataset_list[c].noise_map.real),
+        jnp.array(dataset_list[c].noise_map.imag),
+        tm_real, tm_imag,
+        jnp.asarray(inv_c.reconstruction),
+        jnp.asarray(inv_c.curvature_matrix),
+        jnp.array(inv_c.regularization_matrix),
+        jnp.array(np.asarray(inv_c.mapper_indices)),
+    )
+    log_evidence_check_per_channel.append(float(le_c))
+
+cube_log_evidence_check = float(sum(log_evidence_check_per_channel))
+print(
+    f"\n  cube log_evidence (per-step recompute) = {cube_log_evidence_check:.6f}"
+)
+print(f"  cube log_evidence (reference)          = {cube_log_evidence_ref:.6f}")
+
+np.testing.assert_allclose(
+    cube_log_evidence_check,
+    cube_log_evidence_ref,
+    rtol=1e-4,
+    err_msg=(
+        "Per-step cube log_evidence does not match summed FitInterferometer.log_evidence"
+    ),
+)
+print(
+    "  Assertion PASSED: per-step cube log_evidence matches summed "
+    "FitInterferometer.log_evidence at rtol=1e-4"
+)
+
+
+# ===================================================================
+# PART C — Full-pipeline cube JIT (sum of per-channel log_likelihoods)
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("FULL-PIPELINE CUBE JIT (for comparison)")
+print("=" * 70)
+
+analysis_list = [
+    al.AnalysisInterferometer(dataset=d, adapt_images=adapt_images, use_jax=True)
+    for d in dataset_list
+]
+
+
+def full_cube_pipeline_from_params(params_tree):
+    """Cube log-evidence via the explicit per-channel sum.
+
+    Same shape as the user-facing ``datacube/likelihood_function.py``:
+    feeds the shared instance to every per-channel
+    ``AnalysisInterferometer.log_likelihood_function`` and sums.
+    """
+    total = jnp.zeros(())
+    for analysis in analysis_list:
+        total = total + analysis.log_likelihood_function(instance=params_tree)
+    return total
+
+
+_, full_cube_result = jit_profile(
+    full_cube_pipeline_from_params, "full_cube_pipeline", params_tree
+)
+full_pipeline_per_call = timer.records[-1][1] / 10
+
+print(f"  full cube log_evidence (JIT) = {full_cube_result}")
+
+np.testing.assert_allclose(
+    float(full_cube_result),
+    cube_log_evidence_ref,
+    rtol=1e-4,
+    err_msg="Full-pipeline cube JIT log_evidence does not match summed eager FitInterferometer.log_evidence",
+)
+print("  Eager-vs-JIT cube correctness PASSED")
+
+# ===================================================================
+# PART D — vmap (skipped for cube)
+# ===================================================================
+#
+# The natural batching axis for a cube fit is "datasets" (one entry per
+# channel), not "parameters" (which the interferometer-sibling vmap exercises).
+# vmap-over-channels would require a different graph shape and isn't where the
+# shared-Lᵀ W̃ L optimisation lives. Skipped.
+
+print("\n--- vmap (skipped) ---")
+print(
+    "  Cube batching dimension is 'datasets', not 'parameters'. The "
+    "interferometer-sibling vmap pattern doesn't map cleanly here. Skipped."
+)
+
+# ===================================================================
+# Summary + JSON + bar chart
+# ===================================================================
+
+import json
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+al_version = al.__version__
+
+print("\n" + "=" * 70)
+print(f"JAX LIKELIHOOD FUNCTION SUMMARY — CUBE {instrument.upper()} × {n_channels} — v{al_version}")
+print("=" * 70)
+print(f"  Instrument:              {instrument}")
+print(f"  Channels:                {n_channels}")
+print(f"  Pixel scale:             {pixel_scale} arcsec/pixel")
+print(f"  Real-space mask radius:  {mask_radius} arcsec")
+print(f"  Real-space grid shape:   {real_space_shape[0]} x {real_space_shape[1]}")
+print(f"  Visibilities/chan:       {n_visibilities}")
+print(f"  Delaunay vertices:       {n_mesh_vertices}")
+print(f"  Edge zeroed pixels:      {edge_pixels_total}")
+print("-" * 70)
+print(f"  Cube reference log_evidence:  {cube_log_evidence_ref}")
+print(f"  Cube JIT log_evidence:        {float(full_cube_result)}")
+print("-" * 70)
+
+max_label = max(len(label) for label, _ in likelihood_steps)
+step_total = 0.0
+for i, (label, per_call) in enumerate(likelihood_steps, 1):
+    print(f"  {i:>2}. {label:<{max_label}}  {per_call:>12.6f} s")
+    step_total += per_call
+
+# Shared-Lᵀ W̃ L optimisation savings estimate:
+# Moving the curvature matrix from per-channel to shared would save
+# (n_channels - 1) × per-channel curvature matrix cost.
+shared_lwl_savings = (n_channels - 1) * curvature_matrix_per_channel
+
+print("-" * 70)
+print(f"      {'TOTAL (step-by-step cube cost)':<{max_label}}  {step_total:>12.6f} s")
+print(f"      {'Full pipeline cube (single JIT)':<{max_label}}  {full_pipeline_per_call:>12.6f} s")
+print(f"      {f'Shared-Lᵀ W̃ L savings (curvature only, est.)':<{max_label}}  {shared_lwl_savings:>12.6f} s")
+print("=" * 70)
+
+# --- Save results dictionary ---
+
+likelihood_summary = {
+    "autolens_version": al_version,
+    "instrument": instrument,
+    "model": "delaunay",
+    "n_channels": n_channels,
+    "configuration": {
+        "pixel_scale_arcsec": pixel_scale,
+        "mask_radius_arcsec": mask_radius,
+        "real_space_shape": list(real_space_shape),
+        "visibilities_per_channel": int(n_visibilities),
+        "overlay_shape": list(overlay_shape),
+        "edge_n_points": edge_n_points,
+        "delaunay_vertices": int(n_mesh_vertices),
+        "edge_zeroed_pixels": int(edge_pixels_total),
+        "regularization_coefficient": regularization_coefficient,
+    },
+    "cube_log_evidence_eager": cube_log_evidence_ref,
+    "cube_log_evidence_jit": float(full_cube_result),
+    "log_evidence_per_channel_eager": [float(le) for le in log_evidence_per_channel],
+    "steps_cube_cost": {label: per_call for label, per_call in likelihood_steps},
+    "per_channel_costs": {
+        "inversion_setup": inversion_setup_per_channel,
+        "data_vector": data_vector_per_channel,
+        "curvature_matrix": curvature_matrix_per_channel,
+        "reconstruction": reconstruction_per_channel,
+        "log_evidence": log_evidence_per_channel_cost,
+    },
+    "total_step_by_step_cube": step_total,
+    "full_pipeline_cube_single_jit": full_pipeline_per_call,
+    "shared_lwl_savings_estimate": shared_lwl_savings,
+    "vmap": "SKIPPED — cube batching axis is 'datasets', not 'parameters'",
+}
+
+results_dir = _workspace_root / "results" / "likelihood" / "datacube"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+dict_path = results_dir / f"delaunay_likelihood_summary_{instrument}_v{al_version}.json"
+dict_path.write_text(json.dumps(likelihood_summary, indent=2))
+print(f"\n  Results dict saved to: {dict_path}")
+
+# --- Save bar chart ---
+
+labels = [label for label, _ in likelihood_steps]
+times = [per_call for _, per_call in likelihood_steps]
+
+fig, ax = plt.subplots(figsize=(11, 6))
+y_pos = range(len(labels))
+# Different colours for shared vs per-channel
+colors = ["#55A868" if "(shared)" in label else "#4C72B0" for label in labels]
+bars = ax.barh(y_pos, times, color=colors, edgecolor="white", height=0.6)
+
+for bar, t in zip(bars, times):
+    ax.text(
+        bar.get_width() + max(times) * 0.01,
+        bar.get_y() + bar.get_height() / 2,
+        f"{t:.6f} s",
+        va="center",
+        fontsize=9,
+    )
+
+ax.axvline(
+    full_pipeline_per_call,
+    color="#C44E52",
+    linestyle="--",
+    linewidth=1.5,
+    label=f"Full pipeline cube (single JIT): {full_pipeline_per_call:.6f} s",
+)
+ax.axvline(
+    shared_lwl_savings,
+    color="#8172B2",
+    linestyle=":",
+    linewidth=1.5,
+    label=f"Shared-Lᵀ W̃ L savings est.: {shared_lwl_savings:.6f} s",
+)
+
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=10)
+ax.invert_yaxis()
+ax.set_xlabel("Cube cost per call (s)", fontsize=11)
+fig.suptitle(
+    f"Delaunay Datacube Likelihood — {instrument.upper()} × {n_channels} channels",
+    fontsize=12,
+    fontweight="bold",
+)
+ax.set_title(
+    f"AutoLens v{al_version}  |  {pixel_scale}\"/px  |  "
+    f"{real_space_shape[0]}x{real_space_shape[1]} real-space  |  "
+    f"{n_visibilities} visibilities/chan  |  {n_mesh_vertices} Delaunay verts  |  "
+    f"step-by-step total: {step_total:.6f} s",
+    fontsize=9,
+)
+ax.legend(loc="lower right", fontsize=9)
+ax.margins(x=0.18)
+fig.tight_layout()
+
+chart_path = results_dir / f"delaunay_likelihood_summary_{instrument}_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to:    {chart_path}")
+
+
+# ===================================================================
+# Regression assertion — deterministic cube log-evidence
+# ===================================================================
+#
+# Identical channels = exact N × single-channel log-evidence.
+EXPECTED_LOG_EVIDENCE_CUBE_SMA = n_channels * -3167.5258928840763
+
+np.testing.assert_allclose(
+    cube_log_evidence_ref,
+    EXPECTED_LOG_EVIDENCE_CUBE_SMA,
+    rtol=1e-4,
+    err_msg=(
+        f"datacube/delaunay[{instrument}]: regression — eager cube log_evidence "
+        f"drifted (got {cube_log_evidence_ref}, expected {EXPECTED_LOG_EVIDENCE_CUBE_SMA})"
+    ),
+)
+print(
+    f"\n  Eager cube regression assertion PASSED: log_evidence matches "
+    f"{EXPECTED_LOG_EVIDENCE_CUBE_SMA:.6f}"
+)
+np.testing.assert_allclose(
+    float(full_cube_result),
+    EXPECTED_LOG_EVIDENCE_CUBE_SMA,
+    rtol=1e-4,
+    err_msg=f"datacube/delaunay[{instrument}]: regression — full cube log_evidence drifted",
+)
+print(f"  Full-pipeline cube regression assertion PASSED")
diff --git a/likelihood/imaging/README.md b/likelihood/imaging/README.md
new file mode 100644
index 0000000..8af3c3c
--- /dev/null
+++ b/likelihood/imaging/README.md
@@ -0,0 +1,57 @@
+# likelihood/imaging
+
+JAX JIT profiling for the PyAutoLens **imaging** likelihood function — strong-lens reconstruction from a single 2D image with a known PSF.
+
+The imaging likelihood path is the most fully JIT-decomposed of the four data types. Each script here profiles the **whole likelihood** as one JIT'd program (the cost a sampler actually pays) **and** every internal step under its own JIT (the per-step breakdown that tells you where the time goes). Three model compositions are covered, all with the same step structure but different source representations:
+
+- **MGE** — multi-Gaussian expansion source. Linear light profiles only, so no profile-subtracted image stage. Simplest source path.
+- **Pixelization** — rectangular pixelized source. Adds mesh-to-image mapping and a regularization matrix.
+- **Delaunay** — Delaunay-triangulated source pixelization with adaptive resolution. Heaviest source path; closest to production reconstructions for irregular-morphology sources.
+
+## Step-by-step structure each script profiles
+
+Per the docstrings in each script, the imaging likelihood decomposes into:
+
+1. Instance from parameter vector
+2. Build `Tracer`
+3. Ray-trace grids through the lens
+4. Compute mapping matrix (per-profile images before PSF)
+5. Compute blurred mapping matrix (PSF convolution)
+6. Compute data vector D
+7. Compute curvature matrix F
+8. Reconstruction via positive-only NNLS
+9. Map reconstruction back to image plane
+10. Chi-squared and log-likelihood
+
+XLA may fuse these differently when compiled as one program vs separate pieces, so per-step timings are approximate. They are still the right tool for identifying which step dominates.
+
+## Scripts
+
+| Script | Source representation | What it profiles |
+|--------|-----------------------|------------------|
+| [`mge.py`](./mge.py) | MGE (linear light profiles only) | Lightest imaging path. Useful as the baseline against which pixelization / Delaunay are compared. |
+| [`pixelization.py`](./pixelization.py) | Rectangular pixelization | Adds source mesh + regularization. |
+| [`delaunay.py`](./delaunay.py) | Delaunay pixelization | Production-style irregular-source reconstruction. |
+
+## Default dataset
+
+`dataset/imaging/hst/` — an HST-resolution mock (pixel scale 0.05″, 21×21 PSF) committed to this repo. Other instruments (`euclid`, `jwst`, `ao`) can be regenerated via the source-of-truth scripts at `autolens_workspace_developer/jax_profiling/dataset_setup/imaging.py` and copied into `dataset/imaging/<instrument>/`.
+
+## Headline run-times (populated by Phase 4)
+
+| Script | Instrument | CPU | Laptop GPU | A100 |
+|--------|------------|-----|------------|------|
+| `mge.py` | HST | _populated_ | _populated_ | _populated_ |
+| `pixelization.py` | HST | _populated_ | _populated_ | _populated_ |
+| `delaunay.py` | HST | _populated_ | _populated_ | _populated_ |
+
+Numbers are the **steady-state per-call cost** (single-JIT, post-warmup), in milliseconds. Phase 4's dashboard auto-fills this from the latest `*_summary_v<version>.json` artifacts under `results/likelihood/imaging/`.
+
+## Output
+
+Each script writes:
+
+```
+results/likelihood/imaging/<script>_likelihood_summary_<instrument>_v<al.__version__>.json
+results/likelihood/imaging/<script>_likelihood_summary_<instrument>_v<al.__version__>.png
+```
diff --git a/likelihood/imaging/delaunay.py b/likelihood/imaging/delaunay.py
new file mode 100644
index 0000000..fd0892e
--- /dev/null
+++ b/likelihood/imaging/delaunay.py
@@ -0,0 +1,1191 @@
+"""
+JAX Profiling: Delaunay Imaging Likelihood (Step-by-Step)
+=========================================================
+
+Profiles each step of the JAX likelihood function for an imaging dataset where
+the source galaxy is reconstructed using a Delaunay triangulation mesh with
+ConstantSplit regularization.
+
+Key differences from the rectangular pixelization profiling script:
+
+- Mesh vertices are computed in the **image-plane** via an Overlay grid, then
+  ray-traced to the source-plane (rectangular computes directly in source-plane).
+- Edge points are appended around the mask border and zeroed during inversion.
+- Uses **InterpolatorDelaunay** (barycentric interpolation within triangles)
+  instead of bilinear interpolation on a rectangular grid.
+- Uses **ConstantSplit** regularization (cross-derivative scheme) instead of
+  the simpler Constant neighbour-difference scheme.
+- Delaunay triangulation itself uses scipy on CPU and cannot be JIT-compiled.
+
+Pipeline steps:
+
+1. Ray-trace data grid to source plane
+2. Ray-trace mesh grid (image-plane vertices) to source plane
+3. Lens light images (pre-PSF, JIT) + PSF convolution (eager)
+4. Profile-subtracted image
+5. Border relocation (data grid + mesh grid)
+6. Delaunay triangulation + interpolation + mapper
+7. Mapping matrix
+8. Blurred mapping matrix (PSF convolution)
+9. Data vector (D)
+10. Curvature matrix (F)
+11. Regularization matrix (H) — ConstantSplit scheme
+12. Regularized reconstruction: s = (F + H)^{-1} D
+13. Map reconstruction to image + log evidence
+
+Caveat: XLA may fuse operations differently when compiled as one program vs
+separate pieces, so per-step timings are approximate. They are still useful
+for identifying which step dominates.
+
+All JAX timings use `block_until_ready()` to force synchronous measurement.
+"""
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import time
+import subprocess
+import sys
+from pathlib import Path
+from contextlib import contextmanager
+
+import autofit as af
+import autolens as al
+import autoarray as aa
+from autofit.jax import register_model as _register_model_pytrees
+
+# ---------------------------------------------------------------------------
+# Instrument configuration
+# ---------------------------------------------------------------------------
+
+INSTRUMENTS = {
+    "euclid": {"pixel_scale": 0.1},
+    "hst": {"pixel_scale": 0.05},
+    "jwst": {"pixel_scale": 0.03},
+    "ao": {"pixel_scale": 0.01},
+}
+
+instrument = "hst"  # <-- change this to profile a different instrument
+
+
+# ---------------------------------------------------------------------------
+# Profiling helpers
+# ---------------------------------------------------------------------------
+
+
+class Timer:
+    """Accumulates named timing measurements and prints a summary."""
+
+    def __init__(self):
+        self.records: list[tuple[str, float]] = []
+
+    @contextmanager
+    def section(self, label: str):
+        """Context manager that records wall-clock time for *label*."""
+        start = time.perf_counter()
+        yield
+        elapsed = time.perf_counter() - start
+        self.records.append((label, elapsed))
+        print(f"  [{label}] {elapsed:.4f} s")
+
+    def summary(self):
+        print("\n" + "=" * 70)
+        print("PROFILING SUMMARY")
+        print("=" * 70)
+        max_label = max(len(r[0]) for r in self.records)
+        total = 0.0
+        for label, elapsed in self.records:
+            print(f"  {label:<{max_label}}  {elapsed:>10.4f} s")
+            total += elapsed
+        print("-" * 70)
+        print(f"  {'TOTAL':<{max_label}}  {total:>10.4f} s")
+        print("=" * 70)
+
+
+def block(x):
+    """Call block_until_ready if available (JAX arrays)."""
+    if hasattr(x, "block_until_ready"):
+        x.block_until_ready()
+    return x
+
+
+def jit_profile(func, label, *args, n_repeats=10):
+    """JIT-compile *func*, time first call and steady-state average.
+
+    Returns the compiled function and its result.
+    """
+    jitted = jax.jit(func)
+
+    with timer.section(f"{label}_lower"):
+        lowered = jitted.lower(*args)
+
+    with timer.section(f"{label}_compile"):
+        compiled = lowered.compile()
+
+    with timer.section(f"{label}_first_call"):
+        result = compiled(*args)
+        block(result)
+
+    with timer.section(f"{label}_steady_x{n_repeats}"):
+        for _ in range(n_repeats):
+            result = compiled(*args)
+            block(result)
+
+    per_call = timer.records[-1][1] / n_repeats
+    print(f"    -> per-call avg: {per_call:.6f} s")
+    return compiled, result
+
+
+timer = Timer()
+likelihood_steps = []  # (label, per_call_seconds) for the final summary
+
+# ===================================================================
+# PART A — Setup (not JIT-compiled)
+# ===================================================================
+
+# ---------------------------------------------------------------------------
+# 1. Dataset
+# ---------------------------------------------------------------------------
+
+print(f"\n--- Dataset loading & masking [{instrument}] ---")
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[1]
+pixel_scale = INSTRUMENTS[instrument]["pixel_scale"]
+dataset_path = Path("dataset") / "imaging" / instrument
+
+if al.util.dataset.should_simulate(str(dataset_path)):
+    raise FileNotFoundError(
+        f"Input dataset missing at '{dataset_path}'. The autolens_profiling "
+        f"repo mirrors only the curated datasets needed for default smoke "
+        f"runs. To regenerate or extend datasets, use the source-of-truth "
+        f"scripts under autolens_workspace_developer/jax_profiling/dataset_setup/, "
+        f"then copy the result into autolens_profiling/dataset/."
+    )
+
+with timer.section("dataset_load"):
+    dataset = al.Imaging.from_fits(
+        data_path=dataset_path / "data.fits",
+        psf_path=dataset_path / "psf.fits",
+        noise_map_path=dataset_path / "noise_map.fits",
+        pixel_scales=pixel_scale,
+    )
+
+with timer.section("mask_and_oversample"):
+    mask_radius = 3.5
+
+    mask = al.Mask2D.circular(
+        shape_native=dataset.shape_native,
+        pixel_scales=dataset.pixel_scales,
+        radius=mask_radius,
+    )
+
+    dataset = dataset.apply_mask(mask=mask)
+    dataset = dataset.apply_over_sampling(
+        over_sample_size_lp=4,
+        over_sample_size_pixelization=1,
+    )
+
+    over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(
+        grid=dataset.grid,
+        sub_size_list=[4, 2, 1],
+        radial_list=[0.3, 0.6],
+        centre_list=[(0.0, 0.0)],
+    )
+
+    dataset = dataset.apply_over_sampling(
+        over_sample_size_lp=over_sample_size,
+        over_sample_size_pixelization=1,
+    )
+
+# ---------------------------------------------------------------------------
+# 2. Image mesh + edge points (Delaunay-specific)
+# ---------------------------------------------------------------------------
+
+print("\n--- Image mesh construction (Delaunay) ---")
+
+overlay_shape = (
+    39,
+    39,
+)  # calibrated → 1231 mesh vertices (1201 inside + 30 edge), science fiducial near 1250
+edge_n_points = 30
+
+with timer.section("image_mesh_overlay"):
+    image_mesh = al.image_mesh.Overlay(shape=overlay_shape)
+    image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=dataset.mask)
+
+with timer.section("edge_points"):
+    edge_pixels_total = image_plane_mesh_grid.shape[0]
+    image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(
+        image_plane_mesh_grid=image_plane_mesh_grid,
+        centre=(0.0, 0.0),
+        radius=mask_radius,
+        n_points=edge_n_points,
+    )
+    edge_pixels_total = image_plane_mesh_grid.shape[0] - edge_pixels_total
+
+n_mesh_vertices = image_plane_mesh_grid.shape[0]
+print(f"  Overlay shape: {overlay_shape}")
+print(f"  Mesh vertices (incl. edge): {n_mesh_vertices}")
+print(f"  Edge points added: {edge_pixels_total}")
+
+# ---------------------------------------------------------------------------
+# 3. Model construction
+# ---------------------------------------------------------------------------
+
+print("\n--- Model construction ---")
+
+with timer.section("model_build"):
+    # GaussianPrior(mean=truth, sigma=small) centres prior-median at the
+    # simulator truth while keeping params free so gradient diagnostics
+    # have dimensionality.
+    # Lens light: MGE-60 (full production-fiducial) — replaces single Sersic.
+    # The 60 linear Gaussians enter the inversion's mapping matrix
+    # alongside the source-pixel columns.
+    lens_bulge = al.model_util.mge_model_from(
+        mask_radius=mask_radius,
+        total_gaussians=60,
+        centre_prior_is_uniform=True,
+    )
+
+    mass = af.Model(al.mp.Isothermal)
+    mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.005)
+    mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.005)
+    mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05)
+    _lens_mass_ell = al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0)
+    mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=_lens_mass_ell[0], sigma=0.01)
+    mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=_lens_mass_ell[1], sigma=0.01)
+
+    shear = af.Model(al.mp.ExternalShear)
+    shear.gamma_1 = af.GaussianPrior(mean=0.05, sigma=0.005)
+    shear.gamma_2 = af.GaussianPrior(mean=0.05, sigma=0.005)
+
+    lens = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear)
+
+    mesh = al.mesh.Delaunay(
+        pixels=n_mesh_vertices,
+        zeroed_pixels=edge_pixels_total,
+    )
+    regularization = al.reg.ConstantSplit(coefficient=1.0)
+    pixelization = al.Pixelization(mesh=mesh, regularization=regularization)
+
+    source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)
+
+    model = af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+print(f"  Total free parameters: {model.total_free_parameters}")
+print(f"  Delaunay pixels: {n_mesh_vertices}")
+print(f"  Zeroed edge pixels: {edge_pixels_total}")
+
+# ---------------------------------------------------------------------------
+# 4. Instantiate concrete objects from prior medians
+# ---------------------------------------------------------------------------
+
+print("\n--- Instantiate concrete model ---")
+
+with timer.section("instance_from_vector"):
+    param_vector = model.physical_values_from_prior_medians
+    instance = model.instance_from_vector(vector=param_vector)
+
+with timer.section("register_pytrees"):
+    _register_model_pytrees(model)
+
+params_tree = jax.tree_util.tree_map(jnp.asarray, instance)
+
+n_pytree_leaves = len(jax.tree_util.tree_leaves(params_tree))
+print(f"  Pytree JAX leaves: {n_pytree_leaves}")
+
+tracer = al.Tracer(galaxies=list(instance.galaxies))
+
+# AdaptImages tells FitImaging where mesh vertices live in image-plane
+adapt_images = al.AdaptImages(
+    galaxy_image_plane_mesh_grid_dict={
+        instance.galaxies.source: image_plane_mesh_grid,
+    },
+    galaxy_name_image_plane_mesh_grid_dict={
+        "('galaxies', 'source')": image_plane_mesh_grid,
+    },
+)
+
+print(f"  Tracer planes: {tracer.total_planes}")
+
+# ---------------------------------------------------------------------------
+# Key configuration that dictates run time
+# ---------------------------------------------------------------------------
+
+n_image_pixels = dataset.data.shape[0]
+n_over_sampled_pixels = dataset.grids.lp.over_sampled.shape[0]
+n_source_pixels = n_mesh_vertices
+
+print("\n--- Configuration (determines run time) ---")
+print(f"  Instrument:              {instrument}")
+print(f"  Pixel scale:             {pixel_scale} arcsec/pixel")
+print(f"  Mask radius:             {mask_radius} arcsec")
+print(f"  Image pixels (masked):   {n_image_pixels}")
+print(f"  Over-sampled pixels:     {n_over_sampled_pixels}")
+print(f"  Delaunay vertices:       {n_source_pixels}")
+print(f"  Edge zeroed pixels:      {edge_pixels_total}")
+
+# ---------------------------------------------------------------------------
+# 5. Full-pipeline reference (FitImaging) — eager baseline
+# ---------------------------------------------------------------------------
+
+print("\n--- Full FitImaging (eager baseline) ---")
+
+with timer.section("fit_imaging_eager"):
+    fit = al.FitImaging(
+        dataset=dataset,
+        tracer=tracer,
+        adapt_images=adapt_images,
+        settings=al.Settings(use_border_relocator=True),
+        xp=np,
+    )
+    log_evidence_ref = fit.figure_of_merit
+    log_likelihood_ref = fit.log_likelihood
+
+print(f"  figure_of_merit (log_evidence) = {log_evidence_ref}")
+print(f"  log_likelihood                 = {log_likelihood_ref}")
+
+
+# ===================================================================
+# PART B — Per-step JIT profiling
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("PER-STEP JIT PROFILING")
+print("=" * 70)
+
+# Extract raw arrays from autoarray types via .array so they can cross
+# JIT boundaries.  See CLAUDE.md for rationale.
+
+grid_lp_raw = jnp.array(dataset.grids.lp.array)
+grid_pix_raw = jnp.array(dataset.grids.pixelization.array)
+grid_blurring_raw = jnp.array(dataset.grids.blurring.array)
+mesh_grid_raw = jnp.array(image_plane_mesh_grid.array)
+data_array = jnp.array(dataset.data.array)
+noise_map_array = jnp.array(dataset.noise_map.array)
+
+# Keep autoarray objects for eager calls that need them.
+grid_lp = dataset.grids.lp
+grid_blurring = dataset.grids.blurring
+
+# ---------------------------------------------------------------------------
+# Step 1: Ray-trace data grid to source plane
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 1: Ray-trace data grid ---")
+
+with timer.section("ray_trace_data_eager"):
+    traced_grids = tracer.traced_grid_2d_list_from(
+        grid=dataset.grids.pixelization, xp=jnp
+    )
+    for tg in traced_grids:
+        block(tg)
+
+print(f"  Number of planes traced: {len(traced_grids)}")
+
+
+def ray_trace_data_raw(grid_raw):
+    """Wraps ray-tracing so inputs/outputs are raw arrays."""
+    grid = aa.Grid2DIrregular(values=grid_raw, xp=jnp)
+    traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp)
+    return jnp.stack([tg.array for tg in traced])
+
+
+_, traced_data_grids_raw = jit_profile(
+    ray_trace_data_raw, "ray_trace_data_jit", grid_pix_raw
+)
+likelihood_steps.append(("Ray-trace data grid", timer.records[-1][1] / 10))
+
+print(f"  traced_data_grids shape: {traced_data_grids_raw.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 2: Ray-trace mesh grid (image-plane vertices) to source plane
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 2: Ray-trace mesh grid ---")
+
+with timer.section("ray_trace_mesh_eager"):
+    traced_mesh = tracer.traced_grid_2d_list_from(
+        grid=al.Grid2DIrregular(image_plane_mesh_grid), xp=jnp
+    )
+    for tg in traced_mesh:
+        block(tg)
+
+
+def ray_trace_mesh_raw(mesh_raw):
+    """Ray-trace image-plane mesh vertices to source plane."""
+    grid = aa.Grid2DIrregular(values=mesh_raw, xp=jnp)
+    traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp)
+    return jnp.stack([tg.array for tg in traced])
+
+
+_, traced_mesh_grids_raw = jit_profile(
+    ray_trace_mesh_raw, "ray_trace_mesh_jit", mesh_grid_raw
+)
+likelihood_steps.append(("Ray-trace mesh grid", timer.records[-1][1] / 10))
+
+print(f"  traced_mesh_grids shape: {traced_mesh_grids_raw.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 3: Blurred image of non-linear light profiles (lens light)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 3: Blurred image (lens light profiles) ---")
+
+
+# Sub-step 3a: Compute raw lens light images (JIT-profiled)
+def lens_image_raw(grid_raw, blurring_grid_raw):
+    """Compute lens light images on masked + blurring grids (no PSF)."""
+    grid = aa.Grid2DIrregular(values=grid_raw, xp=jnp)
+    blurring_grid = aa.Grid2DIrregular(values=blurring_grid_raw, xp=jnp)
+    image = tracer.image_2d_from(grid=grid, xp=jnp)
+    blurring_image = tracer.image_2d_from(grid=blurring_grid, xp=jnp)
+    return image.array, blurring_image.array
+
+
+with timer.section("lens_image_eager"):
+    img_eager, blur_img_eager = lens_image_raw(grid_lp_raw, grid_blurring_raw)
+    block(img_eager)
+    block(blur_img_eager)
+
+_, (img_jit, blur_img_jit) = jit_profile(
+    lens_image_raw, "lens_image_jit", grid_lp_raw, grid_blurring_raw
+)
+likelihood_steps.append(("Lens light images (pre-PSF)", timer.records[-1][1] / 10))
+
+# Sub-step 3b: PSF convolution
+with timer.section("blurred_image_eager"):
+    blurred_image = tracer.blurred_image_2d_from(
+        grid=grid_lp,
+        psf=dataset.psf,
+        blurring_grid=grid_blurring,
+        xp=jnp,
+    )
+    block(blurred_image)
+
+print(f"  blurred_image shape: {blurred_image.array.shape}")
+
+
+def blurred_image_from_params(params_tree):
+    """Compute blurred image directly from a pytree ModelInstance — fully JIT-traceable."""
+    t = al.Tracer(galaxies=list(params_tree.galaxies))
+    result = t.blurred_image_2d_from(
+        grid=grid_lp,
+        psf=dataset.psf,
+        blurring_grid=grid_blurring,
+        xp=jnp,
+    )
+    return result.array
+
+
+_, blurred_img_jit = jit_profile(
+    blurred_image_from_params, "blurred_image_jit", params_tree
+)
+likelihood_steps.append(("Blurred image (PSF convolution)", timer.records[-1][1] / 10))
+
+# ---------------------------------------------------------------------------
+# Step 4: Profile-subtracted image (lens light subtraction)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 4: Profile-subtracted image ---")
+
+
+def profile_subtract(data, blurred_image):
+    return data - blurred_image
+
+
+with timer.section("profile_subtract_eager"):
+    blurred_img_jnp = jnp.array(blurred_image.array)
+    profile_subtracted = profile_subtract(data_array, blurred_img_jnp)
+    block(profile_subtracted)
+
+_, profile_subtracted = jit_profile(
+    profile_subtract, "profile_subtract_jit", data_array, blurred_img_jnp
+)
+likelihood_steps.append(("Profile-subtracted image", timer.records[-1][1] / 10))
+
+print(f"  profile_subtracted shape: {profile_subtracted.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 5: Border relocation (data grid + mesh grid)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 5: Border relocation ---")
+
+from autoarray.inversion.mesh.border_relocator import BorderRelocator
+
+with timer.section("border_relocator_setup"):
+    border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1)
+
+traced_source_grid = tracer.traced_grid_2d_list_from(
+    grid=dataset.grids.pixelization, xp=jnp
+)[-1]
+traced_mesh_source = tracer.traced_grid_2d_list_from(
+    grid=al.Grid2DIrregular(image_plane_mesh_grid), xp=jnp
+)[-1]
+
+with timer.section("border_relocation_eager"):
+    relocated_grid = border_relocator.relocated_grid_from(grid=traced_source_grid)
+    relocated_mesh_grid = border_relocator.relocated_mesh_grid_from(
+        grid=traced_source_grid,
+        mesh_grid=traced_mesh_source,
+    )
+    block(relocated_grid)
+    block(relocated_mesh_grid)
+
+print(f"  relocated_data_grid shape: {relocated_grid.array.shape}")
+print(f"  relocated_mesh_grid shape: {relocated_mesh_grid.array.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 6: Delaunay triangulation + interpolation + mapper
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 6: Delaunay triangulation + Interpolation + Mapper ---")
+
+pixelization_obj = instance.galaxies.source.pixelization
+
+with timer.section("delaunay_interpolation_and_mapper"):
+    interpolator = al.InterpolatorDelaunay(
+        mesh=pixelization_obj.mesh,
+        mesh_grid=relocated_mesh_grid,
+        data_grid=relocated_grid,
+    )
+    mapper = al.Mapper(
+        interpolator=interpolator,
+        image_plane_mesh_grid=image_plane_mesh_grid,
+        xp=jnp,
+    )
+
+print(f"  mapper.pixels (source): {mapper.pixels}")
+print(f"  pix_indexes shape: {mapper.pix_indexes_for_sub_slim_index.shape}")
+
+# ---------------------------------------------------------------------------
+# Steps 7-13: Extract matrices from FitImaging inversion for consistency
+# ---------------------------------------------------------------------------
+# The FitImaging pipeline handles edge pixel zeroing, curvature diagonal
+# adjustments, and settings that are difficult to replicate manually.
+# We extract the correct matrices from fit.inversion so the step-by-step
+# matches the reference, then JIT-profile the linear algebra operations.
+
+print("\n--- Extracting inversion matrices from FitImaging ---")
+
+inversion = fit.inversion
+
+with timer.section("extract_inversion_matrices"):
+    bmm_ref = jnp.array(inversion.operated_mapping_matrix)
+    mapping_matrix_ref = jnp.array(inversion.mapping_matrix)
+
+    inv_mapper = inversion.cls_list_from(cls=al.Mapper)[0]
+    neighbors = inv_mapper.neighbors
+    neighbors_array = jnp.array(np.asarray(neighbors))
+    neighbors_sizes = jnp.array(neighbors.sizes)
+
+print(f"  operated_mapping_matrix shape: {bmm_ref.shape}")
+print(f"  mapping_matrix shape: {mapping_matrix_ref.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 7: Mapping matrix
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 7: Mapping matrix ---")
+
+with timer.section("mapping_matrix"):
+    mapping_matrix = inv_mapper.mapping_matrix
+
+print(f"  mapping_matrix shape: {mapping_matrix.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 8: Blurred mapping matrix (PSF convolution)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 8: Blurred mapping matrix ---")
+
+with timer.section("blurred_mapping_matrix"):
+    blurred_mapping_matrix = dataset.psf.convolved_mapping_matrix_from(
+        mapping_matrix=mapping_matrix,
+        mask=dataset.mask,
+        xp=jnp,
+    )
+    block(blurred_mapping_matrix)
+
+# JIT-profile the full inversion setup pipeline (steps 5-8 combined):
+# border relocation → Delaunay triangulation → interpolation → mapper → mapping matrix → PSF convolution.
+# These steps are tightly sequential; the full pipeline JIT-compiles them all together.
+
+
+def blurred_mm_from_params(params_tree):
+    """Compute blurred mapping matrix via full inversion setup from a pytree ModelInstance."""
+    t = al.Tracer(galaxies=list(params_tree.galaxies))
+    # Recreate adapt_images with new galaxy instance so dict lookup by object identity works.
+    adapt_images_jax = al.AdaptImages(
+        galaxy_image_plane_mesh_grid_dict={
+            params_tree.galaxies.source: image_plane_mesh_grid,
+        },
+        galaxy_name_image_plane_mesh_grid_dict={
+            "('galaxies', 'source')": image_plane_mesh_grid,
+        },
+    )
+    fit_jax = al.FitImaging(
+        dataset=dataset,
+        tracer=t,
+        adapt_images=adapt_images_jax,
+        settings=al.Settings(use_border_relocator=True),
+        xp=jnp,
+    )
+    return jnp.array(fit_jax.inversion.operated_mapping_matrix)
+
+
+_, bmm_jit = jit_profile(blurred_mm_from_params, "inversion_setup_jit", params_tree)
+likelihood_steps.append(
+    ("Inversion setup (steps 5-8 combined)", timer.records[-1][1] / 10)
+)
+
+print(f"  blurred_mapping_matrix (JIT) shape: {bmm_jit.shape}")
+
+bmm_jnp = bmm_ref  # Use the reference matrices for linear algebra steps
+print(f"  blurred_mapping_matrix shape: {blurred_mapping_matrix.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 9: Data vector (D)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 9: Data vector ---")
+
+
+def compute_data_vector(blurred_mapping_matrix, image, noise_map):
+    return al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from(
+        blurred_mapping_matrix=blurred_mapping_matrix,
+        image=image,
+        noise_map=noise_map,
+    )
+
+
+profile_sub_jnp = jnp.array(fit.profile_subtracted_image.array)
+noise_jnp = jnp.array(dataset.noise_map.array)
+
+with timer.section("data_vector_eager"):
+    data_vector = compute_data_vector(bmm_jnp, profile_sub_jnp, noise_jnp)
+    block(data_vector)
+
+_, data_vector = jit_profile(
+    compute_data_vector, "data_vector_jit", bmm_jnp, profile_sub_jnp, noise_jnp
+)
+likelihood_steps.append(("Data vector (D)", timer.records[-1][1] / 10))
+
+print(f"  data_vector shape: {data_vector.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 10: Curvature matrix (F)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 10: Curvature matrix ---")
+
+no_reg_list = list(inversion.no_regularization_index_list)
+
+
+def compute_curvature_matrix(blurred_mapping_matrix, noise_map):
+    return al.util.inversion.curvature_matrix_via_mapping_matrix_from(
+        mapping_matrix=blurred_mapping_matrix,
+        noise_map=noise_map,
+        settings=fit.settings,
+        add_to_curvature_diag=True,
+        no_regularization_index_list=no_reg_list,
+        xp=jnp,
+    )
+
+
+with timer.section("curvature_matrix_eager"):
+    curvature_matrix = compute_curvature_matrix(bmm_jnp, noise_jnp)
+    block(curvature_matrix)
+
+_, curvature_matrix = jit_profile(
+    compute_curvature_matrix, "curvature_matrix_jit", bmm_jnp, noise_jnp
+)
+likelihood_steps.append(("Curvature matrix (F)", timer.records[-1][1] / 10))
+
+print(f"  curvature_matrix shape: {curvature_matrix.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 11: Regularization matrix (H) — ConstantSplit scheme
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 11: Regularization matrix (ConstantSplit) ---")
+
+# ConstantSplit uses a cross-derivative scheme via the interpolator's
+# _mappings_sizes_weights_split, not the simple neighbour-difference approach.
+# We extract it from the inversion for consistency and JIT-profile separately.
+
+with timer.section("regularization_matrix_eager"):
+    regularization_matrix = jnp.array(inversion.regularization_matrix)
+    block(regularization_matrix)
+
+likelihood_steps.append(("Regularization matrix (H)", timer.records[-1][1]))
+
+print(f"  regularization_matrix shape: {regularization_matrix.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 12: Regularized reconstruction: s = NNLS(F + H, D)
+# ---------------------------------------------------------------------------
+#
+# Uses ``reconstruction_positive_only_from`` (NNLS) to match production
+# AnalysisImaging behaviour. An earlier version of this script used
+# ``jnp.linalg.solve(F+H, D)`` which under-reports the per-step
+# reconstruction cost (~5 ms vs ~36 ms NNLS on RTX 2060). The two
+# solvers happen to produce identical reconstructions for the
+# well-conditioned ConstantSplit setup at prior medians (no negative
+# source pixels, NNLS reduces to linear solve), so the downstream
+# log-evidence value is unchanged within rtol=1e-4.
+
+print("\n--- Step 12: Regularized reconstruction ---")
+
+
+def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix):
+    curvature_reg_matrix = curvature_matrix + regularization_matrix
+    return al.util.inversion.reconstruction_positive_only_from(
+        data_vector=data_vector,
+        curvature_reg_matrix=curvature_reg_matrix,
+        xp=jnp,
+    )
+
+
+with timer.section("reconstruction_eager"):
+    reconstruction = compute_reconstruction(
+        jnp.array(data_vector),
+        jnp.array(curvature_matrix),
+        jnp.array(regularization_matrix),
+    )
+    block(reconstruction)
+
+_, reconstruction = jit_profile(
+    compute_reconstruction,
+    "reconstruction_jit",
+    jnp.array(data_vector),
+    jnp.array(curvature_matrix),
+    jnp.array(regularization_matrix),
+)
+likelihood_steps.append(("Regularized reconstruction", timer.records[-1][1] / 10))
+
+print(f"  reconstruction shape: {reconstruction.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 13: Map reconstruction to image + log evidence
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 13: Mapped reconstruction + log evidence ---")
+
+
+def compute_log_evidence(
+    data,
+    noise_map,
+    blurred_image,
+    blurred_mapping_matrix,
+    reconstruction,
+    curvature_matrix,
+    regularization_matrix,
+    mapper_indices,
+):
+    """Compute the full log evidence including all five terms:
+
+    -2 ln e = chi^2 + s^T H s + ln[det(F+H)] - ln[det(H)] + noise_norm
+
+    Matches the production formula in
+    ``autoarray/inversion/inversion/abstract.py:log_det_*`` —
+    reduces both the curvature_reg_matrix and the regularization_matrix
+    to the rows/cols indexed by ``mapper_indices`` before the log_det.
+    This drops the no-regularization rows (e.g. MGE Basis linear
+    components) which otherwise make ``det(H) = 0`` and the slogdet
+    return -inf, then uses Cholesky for a numerically stable log_det.
+    """
+    # Map reconstruction to image
+    mapped_recon = al.util.inversion.mapped_reconstructed_data_via_mapping_matrix_from(
+        mapping_matrix=blurred_mapping_matrix,
+        reconstruction=reconstruction,
+        xp=jnp,
+    )
+
+    # model_data = lens light + pixelized source
+    model_data = blurred_image + mapped_recon
+
+    # Chi-squared
+    residual = data - model_data
+    chi_squared = jnp.sum((residual / noise_map) ** 2)
+
+    # Regularization term: s^T H s
+    regularization_term = jnp.dot(
+        reconstruction, jnp.dot(regularization_matrix, reconstruction)
+    )
+
+    # Curvature + regularization matrix
+    curvature_reg_matrix = curvature_matrix + regularization_matrix
+
+    # Reduce to the pixelization rows/cols only, matching the production
+    # ``regularization_matrix_reduced`` / ``curvature_reg_matrix_reduced``
+    # properties — necessary for models with non-regularized linear
+    # components (e.g. MGE lens light).
+    creg_reduced = curvature_reg_matrix[mapper_indices][:, mapper_indices]
+    reg_reduced = regularization_matrix[mapper_indices][:, mapper_indices]
+
+    # Cholesky-based log_det (matches production):
+    #   2 * sum(log(diag(L)))  where L is the lower Cholesky factor.
+    log_det_curvature_reg = 2.0 * jnp.sum(
+        jnp.log(jnp.diag(jnp.linalg.cholesky(creg_reduced)))
+    )
+    log_det_regularization = 2.0 * jnp.sum(
+        jnp.log(jnp.diag(jnp.linalg.cholesky(reg_reduced)))
+    )
+
+    # Noise normalization
+    noise_normalization = jnp.sum(jnp.log(2 * jnp.pi * noise_map**2))
+
+    return -0.5 * (
+        chi_squared
+        + regularization_term
+        + log_det_curvature_reg
+        - log_det_regularization
+        + noise_normalization
+    )
+
+
+# For the JIT profiling we use the step-by-step matrices for timing.
+# For the correctness assertion we use the inversion's own matrices, because
+# cumulative floating-point differences between JIT-compiled and eager paths
+# (especially through ill-conditioned solves) can compound significantly.
+
+blurred_img_jnp = jnp.array(blurred_image.array)
+recon_jnp = jnp.array(reconstruction)
+curv_jnp = jnp.array(curvature_matrix)
+reg_jnp = jnp.array(regularization_matrix)
+mapper_indices_jnp = jnp.array(np.asarray(inversion.mapper_indices))
+
+with timer.section("log_evidence_eager"):
+    log_evidence = compute_log_evidence(
+        data_array,
+        noise_jnp,
+        blurred_img_jnp,
+        bmm_jnp,
+        recon_jnp,
+        curv_jnp,
+        reg_jnp,
+        mapper_indices_jnp,
+    )
+    block(log_evidence)
+
+_, log_evidence = jit_profile(
+    compute_log_evidence,
+    "log_evidence_jit",
+    data_array,
+    noise_jnp,
+    blurred_img_jnp,
+    bmm_jnp,
+    recon_jnp,
+    curv_jnp,
+    reg_jnp,
+    mapper_indices_jnp,
+)
+likelihood_steps.append(("Mapped recon + log evidence", timer.records[-1][1] / 10))
+
+print(f"  log_evidence (step-by-step) = {log_evidence}")
+
+# Correctness check: recompute log_evidence using the inversion's own
+# reconstruction and curvature matrix to avoid accumulated FP drift.
+inv_recon_jnp = jnp.array(inversion.reconstruction)
+inv_curv_jnp = jnp.array(inversion.curvature_matrix)
+
+log_evidence_check = compute_log_evidence(
+    data_array,
+    noise_jnp,
+    blurred_img_jnp,
+    bmm_jnp,
+    inv_recon_jnp,
+    inv_curv_jnp,
+    reg_jnp,
+    mapper_indices_jnp,
+)
+print(f"  log_evidence (inv matrices) = {log_evidence_check}")
+print(f"  log_evidence (reference)    = {log_evidence_ref}")
+
+np.testing.assert_allclose(
+    float(log_evidence_check),
+    float(log_evidence_ref),
+    rtol=1e-4,
+    err_msg="Log_evidence from inversion matrices does not match FitImaging.log_evidence",
+)
+print(
+    "  Assertion PASSED: inversion-matrix log_evidence matches FitImaging.log_evidence"
+)
+
+# ===================================================================
+# PART C — Full-pipeline JIT for comparison
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("FULL-PIPELINE JIT (for comparison)")
+print("=" * 70)
+
+analysis = al.AnalysisImaging(dataset=dataset, adapt_images=adapt_images, use_jax=True)
+
+
+def full_pipeline_from_params(params_tree):
+    return analysis.log_likelihood_function(instance=params_tree)
+
+
+_, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree)
+full_pipeline_per_call = timer.records[-1][1] / 10
+
+print(f"  full log_likelihood = {full_result}")
+
+# ===================================================================
+# PART D — vmap + correctness
+# ===================================================================
+
+print("\n--- vmap batched evaluation ---")
+
+# WARNING: The vmap compilation for the Delaunay pipeline takes 20+ minutes on CPU.
+# The XLA graph for a batched Delaunay inversion (including scipy triangulation,
+# border relocation, interpolation, mapping matrix construction, and PSF convolution)
+# is extremely large. The single-call JIT above compiles in ~2s and runs in ~1.8s,
+# but vmap recompiles the entire graph for batch_size independent evaluations.
+#
+# This is likely a candidate for optimisation — either via custom_vjp to avoid
+# retracing the full pipeline, or by restructuring the Delaunay steps to reduce
+# the XLA graph size. For now, skip vmap by default and run it only when explicitly
+# requested via DELAUNAY_VMAP=1 environment variable.
+
+import os
+
+run_vmap = os.environ.get("DELAUNAY_VMAP", "0") == "1"
+
+if not run_vmap:
+    print("  SKIPPED: vmap compilation takes 20+ minutes for Delaunay pipeline.")
+    print("  Set DELAUNAY_VMAP=1 to run this section.")
+    vmap_batch_time = None
+    vmap_per_call = None
+    vmap_speedup = None
+else:
+
+    batch_size = 3
+    parameters = jax.tree_util.tree_map(
+        lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)),
+        params_tree,
+    )
+
+    vmapped_full = jax.jit(jax.vmap(full_pipeline_from_params))
+
+    with timer.section("vmap_first_call"):
+        result_vmap = vmapped_full(parameters)
+        block(result_vmap)
+
+    n_vmap_repeats = 10
+    with timer.section(f"vmap_steady_x{n_vmap_repeats}"):
+        for _ in range(n_vmap_repeats):
+            result_vmap = vmapped_full(parameters)
+            block(result_vmap)
+
+    vmap_batch_time = timer.records[-1][1] / n_vmap_repeats
+    vmap_per_call = vmap_batch_time / batch_size
+    vmap_speedup = full_pipeline_per_call / vmap_per_call
+
+    print(f"  batch results = {result_vmap}")
+    print(f"  vmap batch of {batch_size}:   {vmap_batch_time:.6f} s")
+    print(f"  vmap per call:         {vmap_per_call:.6f} s")
+    print(f"  single JIT per call:   {full_pipeline_per_call:.6f} s")
+    print(f"  vmap speedup:          {vmap_speedup:.1f}x faster per likelihood")
+
+    np.testing.assert_allclose(
+        np.array(result_vmap),
+        float(full_result),
+        rtol=1e-4,
+        err_msg="delaunay: JAX vmap likelihood mismatch",
+    )
+    print("  Correctness check PASSED")
+
+    # --- Static memory analysis ---
+
+    print("\n--- Static memory analysis ---")
+
+    lowered_batched = vmapped_full.lower(parameters)
+    compiled_batched = lowered_batched.compile()
+
+    memory_analysis = compiled_batched.memory_analysis()
+    print(f"  Output size:  {memory_analysis.output_size_in_bytes / 1024**2:.3f} MB")
+    print(f"  Temp size:    {memory_analysis.temp_size_in_bytes / 1024**2:.3f} MB")
+    print(
+        f"  Total:        "
+        f"{(memory_analysis.output_size_in_bytes + memory_analysis.temp_size_in_bytes) / 1024**2:.3f} MB"
+    )
+
+
+# ===================================================================
+# JAX Likelihood Function Summary
+# ===================================================================
+
+import json
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+al_version = al.__version__
+
+print("\n" + "=" * 70)
+print(f"JAX LIKELIHOOD FUNCTION SUMMARY — {instrument.upper()} — v{al_version}")
+print("=" * 70)
+print(f"  Instrument:            {instrument}")
+print(f"  Pixel scale:           {pixel_scale} arcsec/pixel")
+print(f"  Mask radius:           {mask_radius} arcsec")
+print(f"  Image pixels (masked): {n_image_pixels}")
+print(f"  Over-sampled pixels:   {n_over_sampled_pixels}")
+print(f"  Delaunay vertices:     {n_source_pixels}")
+print(f"  Edge zeroed pixels:    {edge_pixels_total}")
+print("-" * 70)
+
+max_label = max(len(label) for label, _ in likelihood_steps)
+step_total = 0.0
+for i, (label, per_call) in enumerate(likelihood_steps, 1):
+    print(f"  {i:>2}. {label:<{max_label}}  {per_call:>12.6f} s")
+    step_total += per_call
+
+print("-" * 70)
+print(f"      {'TOTAL (step-by-step)':<{max_label}}  {step_total:>12.6f} s")
+print(
+    f"      {'Full pipeline (single JIT)':<{max_label}}  {full_pipeline_per_call:>12.6f} s"
+)
+if vmap_per_call is not None:
+    print(f"      {f'vmap batch (per call)':<{max_label}}  {vmap_per_call:>12.6f} s")
+    print(f"      {f'vmap speedup vs single JIT':<{max_label}}  {vmap_speedup:>11.1f}x")
+else:
+    print(f"      {'vmap':<{max_label}}  {'SKIPPED':>12}")
+print("=" * 70)
+
+# --- Save results dictionary ---
+
+likelihood_summary = {
+    "autolens_version": al_version,
+    "instrument": instrument,
+    "configuration": {
+        "pixel_scale_arcsec": pixel_scale,
+        "mask_radius_arcsec": mask_radius,
+        "image_pixels_masked": int(n_image_pixels),
+        "over_sampled_pixels": int(n_over_sampled_pixels),
+        "delaunay_vertices": int(n_source_pixels),
+        "edge_zeroed_pixels": int(edge_pixels_total),
+    },
+    "steps": {label: per_call for label, per_call in likelihood_steps},
+    "total_step_by_step": step_total,
+    "full_pipeline_single_jit": full_pipeline_per_call,
+    "vmap": "SKIPPED — compilation takes 20+ minutes (set DELAUNAY_VMAP=1)",
+}
+
+if vmap_per_call is not None:
+    likelihood_summary["vmap"] = {
+        "batch_size": batch_size,
+        "batch_time": vmap_batch_time,
+        "per_call": vmap_per_call,
+        "speedup_vs_single_jit": round(vmap_speedup, 1),
+    }
+
+results_dir = _workspace_root / "results" / "likelihood" / "imaging"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+dict_path = results_dir / f"delaunay_likelihood_summary_{instrument}_v{al_version}.json"
+dict_path.write_text(json.dumps(likelihood_summary, indent=2))
+print(f"\n  Results dict saved to: {dict_path}")
+
+# --- Save bar chart ---
+
+labels = [label for label, _ in likelihood_steps]
+times = [per_call for _, per_call in likelihood_steps]
+
+fig, ax = plt.subplots(figsize=(10, 6))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color="#4C72B0", edgecolor="white", height=0.6)
+
+for bar, t in zip(bars, times):
+    ax.text(
+        bar.get_width() + max(times) * 0.01,
+        bar.get_y() + bar.get_height() / 2,
+        f"{t:.6f} s",
+        va="center",
+        fontsize=9,
+    )
+
+ax.axvline(
+    full_pipeline_per_call,
+    color="#C44E52",
+    linestyle="--",
+    linewidth=1.5,
+    label=f"Full pipeline (single JIT): {full_pipeline_per_call:.6f} s",
+)
+if vmap_per_call is not None:
+    ax.axvline(
+        vmap_per_call,
+        color="#55A868",
+        linestyle="--",
+        linewidth=1.5,
+        label=f"vmap batch per call: {vmap_per_call:.6f} s ({vmap_speedup:.1f}x faster)",
+    )
+
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=10)
+ax.invert_yaxis()
+ax.set_xlabel("Time per call (s)", fontsize=11)
+fig.suptitle(
+    f"Delaunay Imaging Likelihood — {instrument.upper()}",
+    fontsize=12,
+    fontweight="bold",
+)
+ax.set_title(
+    f'AutoLens v{al_version}  |  {pixel_scale}"/px  |  {n_image_pixels} pixels  |  '
+    f"{n_over_sampled_pixels} over-sampled  |  {n_source_pixels} Delaunay vertices  |  "
+    f"total: {step_total:.6f} s",
+    fontsize=9,
+)
+ax.legend(loc="lower right", fontsize=9)
+ax.margins(x=0.15)
+fig.tight_layout()
+
+chart_path = results_dir / f"delaunay_likelihood_summary_{instrument}_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to:    {chart_path}")
+
+
+# ===================================================================
+# Regression assertion — realistic-scale deterministic log-evidence
+# ===================================================================
+#
+# Simulator truth parameters via GaussianPrior(mean=truth, sigma=small)
+# make the full-pipeline log-evidence deterministic at the prior median.
+# vmap result asserted only when DELAUNAY_VMAP=1 (vmap compile takes 20+ min).
+EXPECTED_LOG_EVIDENCE_HST = 26642.278160003658  # MGE-60 lens light + 39x39 overlay + ConstantSplit (rebaselined 2026-05-12)
+
+np.testing.assert_allclose(
+    log_evidence_ref,
+    EXPECTED_LOG_EVIDENCE_HST,
+    rtol=1e-4,
+    err_msg=(
+        f"imaging/delaunay[{instrument}]: regression — eager log_evidence drifted "
+        f"(got {log_evidence_ref}, expected {EXPECTED_LOG_EVIDENCE_HST})"
+    ),
+)
+print(
+    f"  Eager regression assertion PASSED: log_evidence matches "
+    f"{EXPECTED_LOG_EVIDENCE_HST:.6f}"
+)
+np.testing.assert_allclose(
+    float(full_result),
+    EXPECTED_LOG_EVIDENCE_HST,
+    rtol=1e-4,
+    err_msg=f"imaging/delaunay[{instrument}]: regression — full log_evidence drifted",
+)
+if run_vmap:
+    np.testing.assert_allclose(
+        np.array(result_vmap),
+        EXPECTED_LOG_EVIDENCE_HST,
+        rtol=1e-4,
+        err_msg=f"imaging/delaunay[{instrument}]: regression — vmap log_evidence drifted",
+    )
+print(
+    f"  Regression assertion PASSED: log_evidence matches {EXPECTED_LOG_EVIDENCE_HST:.6f}"
+)
diff --git a/likelihood/imaging/mge.py b/likelihood/imaging/mge.py
new file mode 100644
index 0000000..0f5ceb1
--- /dev/null
+++ b/likelihood/imaging/mge.py
@@ -0,0 +1,878 @@
+"""
+JAX Profiling: MGE Imaging Likelihood (Step-by-Step)
+=====================================================
+
+Profiles each step of the JAX likelihood function for an imaging dataset where
+the lens galaxy's light is modelled with a multi-Gaussian expansion (MGE).
+
+Rather than timing the whole likelihood as a single JIT-compiled block (which
+hides internal bottlenecks), this script JIT-compiles and times each step of
+the pipeline individually:
+
+1. Instance from parameter vector
+2. Build Tracer
+3. Ray-trace grids through the lens
+4. Compute mapping matrix (per-profile images before PSF)
+5. Compute blurred mapping matrix (PSF convolution)
+6. Compute data vector  (D)
+7. Compute curvature matrix  (F)
+8. Reconstruction via positive-only NNLS
+9. Map reconstruction back to image plane
+10. Chi-squared and log likelihood
+
+Note: because the MGE model uses only linear light profiles (lp_linear),
+there is no non-linear blurred image or profile-subtracted image step.
+
+Caveat: XLA may fuse operations differently when compiled as one program vs
+separate pieces, so per-step timings are approximate. They are still useful
+for identifying which step dominates.
+
+All JAX timings use `block_until_ready()` to force synchronous measurement.
+
+Pytree-native parameter inputs (recommended pattern)
+----------------------------------------------------
+
+This script uses ``af.ModelInstance`` as the JIT input via PyAutoFit's
+opt-in pytree registration (``autofit.jax.register_model(model)``). The
+JIT'd closures consume the instance directly, so:
+
+* ``model.instance_from_vector`` is no longer called inside the JIT trace —
+  parameter unpacking happens once at registration time and JAX walks the
+  pytree on every call.
+* Parameter identity is preserved through ``jax.jit`` and ``jax.vmap``;
+  XLA cache keys reflect the structured pytree, not a flat vector shape.
+* ``vmap`` batching is ``jax.tree_util.tree_map`` over the instance leaves
+  — callers no longer have to stack a ``(batch, N)`` array.
+
+New profiling scripts should follow this pattern. The flat-vector path in
+``Fitness.call`` / ``model.instance_from_vector(..., xp=jnp)`` remains the
+production likelihood entry point and is intentionally untouched here.
+"""
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import time
+import subprocess
+import sys
+from pathlib import Path
+from contextlib import contextmanager
+
+import autofit as af
+import autolens as al
+import autoarray as aa
+from autofit.jax import register_model as _register_model_pytrees
+
+# ---------------------------------------------------------------------------
+# Instrument configuration
+# ---------------------------------------------------------------------------
+
+INSTRUMENTS = {
+    "euclid": {"pixel_scale": 0.1},
+    "hst": {"pixel_scale": 0.05},
+    "jwst": {"pixel_scale": 0.03},
+    "ao": {"pixel_scale": 0.01},
+}
+
+instrument = "hst"  # <-- change this to profile a different instrument
+
+
+# ---------------------------------------------------------------------------
+# Profiling helpers
+# ---------------------------------------------------------------------------
+
+
+class Timer:
+    """Accumulates named timing measurements and prints a summary."""
+
+    def __init__(self):
+        self.records: list[tuple[str, float]] = []
+
+    @contextmanager
+    def section(self, label: str):
+        """Context manager that records wall-clock time for *label*."""
+        start = time.perf_counter()
+        yield
+        elapsed = time.perf_counter() - start
+        self.records.append((label, elapsed))
+        print(f"  [{label}] {elapsed:.4f} s")
+
+    def summary(self):
+        print("\n" + "=" * 70)
+        print("PROFILING SUMMARY")
+        print("=" * 70)
+        max_label = max(len(r[0]) for r in self.records)
+        total = 0.0
+        for label, elapsed in self.records:
+            print(f"  {label:<{max_label}}  {elapsed:>10.4f} s")
+            total += elapsed
+        print("-" * 70)
+        print(f"  {'TOTAL':<{max_label}}  {total:>10.4f} s")
+        print("=" * 70)
+
+
+def block(x):
+    """Call block_until_ready if available (JAX arrays)."""
+    if hasattr(x, "block_until_ready"):
+        x.block_until_ready()
+    return x
+
+
+def jit_profile(func, label, *args, n_repeats=10):
+    """JIT-compile *func*, time first call and steady-state average.
+
+    Returns the compiled function and its result.
+    """
+    jitted = jax.jit(func)
+
+    with timer.section(f"{label}_lower"):
+        lowered = jitted.lower(*args)
+
+    with timer.section(f"{label}_compile"):
+        compiled = lowered.compile()
+
+    with timer.section(f"{label}_first_call"):
+        result = compiled(*args)
+        block(result)
+
+    with timer.section(f"{label}_steady_x{n_repeats}"):
+        for _ in range(n_repeats):
+            result = compiled(*args)
+            block(result)
+
+    per_call = timer.records[-1][1] / n_repeats
+    print(f"    -> per-call avg: {per_call:.6f} s")
+    return compiled, result
+
+
+timer = Timer()
+likelihood_steps = []  # (label, per_call_seconds) for the final summary
+
+# ===================================================================
+# PART A — Setup (not JIT-compiled)
+# ===================================================================
+
+# ---------------------------------------------------------------------------
+# 1. Dataset
+# ---------------------------------------------------------------------------
+
+print(f"\n--- Dataset loading & masking [{instrument}] ---")
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[1]
+pixel_scale = INSTRUMENTS[instrument]["pixel_scale"]
+dataset_path = Path("dataset") / "imaging" / instrument
+
+if al.util.dataset.should_simulate(str(dataset_path)):
+    raise FileNotFoundError(
+        f"Input dataset missing at '{dataset_path}'. The autolens_profiling "
+        f"repo mirrors only the curated datasets needed for default smoke "
+        f"runs. To regenerate or extend datasets, use the source-of-truth "
+        f"scripts under autolens_workspace_developer/jax_profiling/dataset_setup/, "
+        f"then copy the result into autolens_profiling/dataset/."
+    )
+
+with timer.section("dataset_load"):
+    dataset = al.Imaging.from_fits(
+        data_path=dataset_path / "data.fits",
+        psf_path=dataset_path / "psf.fits",
+        noise_map_path=dataset_path / "noise_map.fits",
+        pixel_scales=pixel_scale,
+    )
+
+with timer.section("mask_and_oversample"):
+    mask_radius = 3.5
+
+    mask = al.Mask2D.circular(
+        shape_native=dataset.shape_native,
+        pixel_scales=dataset.pixel_scales,
+        radius=mask_radius,
+    )
+
+    dataset = dataset.apply_mask(mask=mask)
+    dataset = dataset.apply_over_sampling(over_sample_size_lp=4)
+
+    over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(
+        grid=dataset.grid,
+        sub_size_list=[4, 2, 1],
+        radial_list=[0.3, 0.6],
+        centre_list=[(0.0, 0.0)],
+    )
+
+    dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)
+
+# ---------------------------------------------------------------------------
+# 2. Model construction
+# ---------------------------------------------------------------------------
+
+print("\n--- Model construction ---")
+
+with timer.section("model_build"):
+    # GaussianPrior(mean=truth, sigma=small) centres prior-median at the
+    # simulator truth while keeping params free so gradient diagnostics
+    # have dimensionality.
+    lens_bulge = al.model_util.mge_model_from(
+        mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True
+    )
+
+    # Mass and shear fixed to simulator truth (not GaussianPrior) because
+    # tracing GaussianPrior-backed mass params through this script's
+    # ``mapping_matrix_from_params`` JIT trigger a pre-existing xp=np/jnp
+    # propagation bug in autogalaxy/profiles/mass/total/isothermal.py:108
+    # (Isothermal.deflections_yx_2d_from called with xp=np on traced inputs).
+    # The bug is specific to this script's MGE-lens-light + over-sampled-LP
+    # combination; the likelihood-only imaging/mge_gradients.py uses the
+    # same pattern without the failing JIT and works under Option A.
+    mass = af.Model(al.mp.Isothermal)
+    mass.centre = (0.0, 0.0)
+    mass.einstein_radius = 1.6
+    mass.ell_comps = al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0)
+
+    shear = af.Model(al.mp.ExternalShear)
+    shear.gamma_1 = 0.05
+    shear.gamma_2 = 0.05
+
+    lens = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear)
+
+    source_bulge = al.model_util.mge_model_from(
+        mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False
+    )
+
+    source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)
+
+    model = af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+print(f"  Total free parameters: {model.total_free_parameters}")
+
+# ---------------------------------------------------------------------------
+# 3. Instantiate concrete objects from prior medians
+# ---------------------------------------------------------------------------
+
+print("\n--- Instantiate concrete model ---")
+
+with timer.section("instance_from_vector"):
+    param_vector = model.physical_values_from_prior_medians
+    instance = model.instance_from_vector(vector=param_vector)
+
+# Register every concrete `model.cls` (Galaxy, profile classes, ModelInstance,
+# Collection, …) with `jax.tree_util` so the instance can cross JIT/vmap
+# boundaries directly. This must happen AFTER the model is built, because
+# registration walks the model's class graph.
+with timer.section("register_pytrees"):
+    _register_model_pytrees(model)
+
+# JIT input: the instance itself, with all parameter leaves promoted to JAX
+# arrays. We keep `instance` (the eager NumPy version) around for any
+# non-JIT setup that needs to read parameter values directly.
+params_tree = jax.tree_util.tree_map(jnp.asarray, instance)
+
+tracer = al.Tracer(galaxies=list(instance.galaxies))
+
+print(f"  Tracer planes: {tracer.total_planes}")
+
+# ---------------------------------------------------------------------------
+# Key configuration that dictates run time
+# ---------------------------------------------------------------------------
+
+n_image_pixels = dataset.data.shape[0]
+n_over_sampled_pixels = dataset.grids.lp.over_sampled.shape[0]
+n_linear_gaussians = len(tracer.cls_list_from(cls=al.lp_linear.LightProfileLinear))
+
+print("\n--- Configuration (determines run time) ---")
+print(f"  Instrument:              {instrument}")
+print(f"  Pixel scale:             {pixel_scale} arcsec/pixel")
+print(f"  Mask radius:             {mask_radius} arcsec")
+print(f"  Image pixels (masked):   {n_image_pixels}")
+print(f"  Over-sampled pixels:     {n_over_sampled_pixels}")
+print(f"  Linear Gaussians:        {n_linear_gaussians}")
+
+# ---------------------------------------------------------------------------
+# 4. Full-pipeline reference (FitImaging) — eager baseline
+# ---------------------------------------------------------------------------
+
+print("\n--- Full FitImaging (eager baseline) ---")
+
+with timer.section("fit_imaging_eager"):
+    fit = al.FitImaging(
+        dataset=dataset,
+        tracer=tracer,
+        settings=al.Settings(use_border_relocator=True),
+        xp=np,
+    )
+    log_evidence_ref = fit.figure_of_merit
+    log_likelihood_ref = fit.log_likelihood
+
+print(f"  figure_of_merit (log_evidence) = {log_evidence_ref}")
+print(f"  log_likelihood                 = {log_likelihood_ref}")
+
+
+# ===================================================================
+# PART B — Per-step JIT profiling
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("PER-STEP JIT PROFILING")
+print("=" * 70)
+
+# Extract raw arrays from autoarray types via .array so they can cross
+# JIT boundaries.  See CLAUDE.md for rationale.
+
+grid_lp_raw = jnp.array(dataset.grids.lp.array)
+data_array = jnp.array(dataset.data.array)
+noise_map_array = jnp.array(dataset.noise_map.array)
+
+# Keep autoarray objects for eager calls that need them.
+grid_lp = dataset.grids.lp
+
+# ---------------------------------------------------------------------------
+# Step 1: Ray-trace grids
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 1: Ray-trace grids ---")
+
+with timer.section("ray_trace_eager"):
+    traced_grids = tracer.traced_grid_2d_list_from(grid=grid_lp, xp=jnp)
+    for tg in traced_grids:
+        block(tg)
+
+print(f"  Number of planes traced: {len(traced_grids)}")
+
+
+def ray_trace_raw(grid_raw):
+    """Wraps ray-tracing so inputs/outputs are raw arrays."""
+    grid = aa.Grid2DIrregular(values=grid_raw, xp=jnp)
+    traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp)
+    return jnp.stack([tg.array for tg in traced])
+
+
+_, traced_grids_raw = jit_profile(ray_trace_raw, "ray_trace_jit", grid_lp_raw)
+likelihood_steps.append(("Ray-trace grids", timer.records[-1][1] / 10))
+
+print(f"  traced_grids shape: {traced_grids_raw.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 2: Build linear objects and mapping matrix
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 2: Mapping matrix (linear profile images) ---")
+
+with timer.section("linear_obj_setup"):
+    tracer_to_inv = al.TracerToInversion(
+        dataset=aa.DatasetInterface(
+            data=fit.profile_subtracted_image,
+            noise_map=dataset.noise_map,
+            grids=dataset.grids,
+            psf=dataset.psf,
+            sparse_operator=dataset.sparse_operator,
+        ),
+        tracer=tracer,
+        settings=al.Settings(use_border_relocator=True),
+    )
+
+    lp_linear_func_galaxy_dict = tracer_to_inv.lp_linear_func_list_galaxy_dict
+
+    lp_linear_funcs = list(lp_linear_func_galaxy_dict.keys())
+
+# mapping_matrix and operated_mapping_matrix_override already return raw arrays.
+with timer.section("mapping_matrix"):
+    mapping_matrices = [func.mapping_matrix for func in lp_linear_funcs]
+    mapping_matrix = (
+        np.hstack(mapping_matrices)
+        if len(mapping_matrices) > 1
+        else mapping_matrices[0]
+    )
+
+print(f"  mapping_matrix shape: {mapping_matrix.shape}")
+
+
+def mapping_matrix_from_params(params_tree):
+    """Compute mapping matrix from a pytree-shaped ``ModelInstance``.
+
+    No flat-vector unpacking inside the trace: ``params_tree`` is the
+    structured instance directly (registered as a JAX pytree above).
+    """
+    t = al.Tracer(galaxies=list(params_tree.galaxies))
+    tti = al.TracerToInversion(
+        dataset=aa.DatasetInterface(
+            data=fit.profile_subtracted_image,
+            noise_map=dataset.noise_map,
+            grids=dataset.grids,
+            psf=dataset.psf,
+            sparse_operator=dataset.sparse_operator,
+        ),
+        tracer=t,
+        settings=al.Settings(use_border_relocator=True),
+        xp=jnp,
+    )
+    funcs = list(tti.lp_linear_func_list_galaxy_dict.keys())
+    matrices = [f.mapping_matrix for f in funcs]
+    return jnp.hstack(matrices) if len(matrices) > 1 else matrices[0]
+
+
+_, mm_jit = jit_profile(mapping_matrix_from_params, "mapping_matrix_jit", params_tree)
+likelihood_steps.append(("Mapping matrix", timer.records[-1][1] / 10))
+
+print(f"  mapping_matrix (JIT) shape: {mm_jit.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 3: Blurred mapping matrix (PSF convolution of each profile)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 3: Blurred mapping matrix ---")
+
+with timer.section("blurred_mapping_matrix"):
+    blurred_matrices = [
+        func.operated_mapping_matrix_override for func in lp_linear_funcs
+    ]
+    blurred_mapping_matrix = (
+        np.hstack(blurred_matrices)
+        if len(blurred_matrices) > 1
+        else blurred_matrices[0]
+    )
+
+print(f"  blurred_mapping_matrix shape: {blurred_mapping_matrix.shape}")
+
+
+def blurred_mm_from_params(params_tree):
+    """Compute blurred mapping matrix from a pytree-shaped ``ModelInstance``."""
+    t = al.Tracer(galaxies=list(params_tree.galaxies))
+    tti = al.TracerToInversion(
+        dataset=aa.DatasetInterface(
+            data=fit.profile_subtracted_image,
+            noise_map=dataset.noise_map,
+            grids=dataset.grids,
+            psf=dataset.psf,
+            sparse_operator=dataset.sparse_operator,
+        ),
+        tracer=t,
+        settings=al.Settings(use_border_relocator=True),
+        xp=jnp,
+    )
+    funcs = list(tti.lp_linear_func_list_galaxy_dict.keys())
+    matrices = [f.operated_mapping_matrix_override for f in funcs]
+    return jnp.hstack(matrices) if len(matrices) > 1 else matrices[0]
+
+
+_, bmm_jit = jit_profile(blurred_mm_from_params, "blurred_mm_jit", params_tree)
+likelihood_steps.append(("Blurred mapping matrix", timer.records[-1][1] / 10))
+
+print(f"  blurred_mapping_matrix (JIT) shape: {bmm_jit.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 4: Data vector (D)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 4: Data vector ---")
+
+
+def compute_data_vector(blurred_mapping_matrix, image, noise_map):
+    return al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from(
+        blurred_mapping_matrix=blurred_mapping_matrix,
+        image=image,
+        noise_map=noise_map,
+    )
+
+
+bmm_jnp = jnp.array(blurred_mapping_matrix)
+noise_jnp = jnp.array(dataset.noise_map.array)
+
+with timer.section("data_vector_eager"):
+    data_vector = compute_data_vector(bmm_jnp, data_array, noise_jnp)
+    block(data_vector)
+
+_, data_vector = jit_profile(
+    compute_data_vector, "data_vector_jit", bmm_jnp, data_array, noise_jnp
+)
+likelihood_steps.append(("Data vector (D)", timer.records[-1][1] / 10))
+
+print(f"  data_vector shape: {data_vector.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 5: Curvature matrix (F)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 5: Curvature matrix ---")
+
+n_linear = bmm_jnp.shape[1]
+
+
+def compute_curvature_matrix(blurred_mapping_matrix, noise_map):
+    return al.util.inversion.curvature_matrix_via_mapping_matrix_from(
+        mapping_matrix=blurred_mapping_matrix,
+        noise_map=noise_map,
+        add_to_curvature_diag=True,
+        no_regularization_index_list=list(range(n_linear)),
+        xp=jnp,
+    )
+
+
+with timer.section("curvature_matrix_eager"):
+    curvature_matrix = compute_curvature_matrix(bmm_jnp, noise_jnp)
+    block(curvature_matrix)
+
+_, curvature_matrix = jit_profile(
+    compute_curvature_matrix, "curvature_matrix_jit", bmm_jnp, noise_jnp
+)
+likelihood_steps.append(("Curvature matrix (F)", timer.records[-1][1] / 10))
+
+print(f"  curvature_matrix shape: {curvature_matrix.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 6: Reconstruction (positive-only NNLS)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 6: Reconstruction (NNLS) ---")
+
+
+def compute_reconstruction(data_vector, curvature_matrix):
+    return al.util.inversion.reconstruction_positive_only_from(
+        data_vector=data_vector,
+        curvature_reg_matrix=curvature_matrix,
+        xp=jnp,
+    )
+
+
+with timer.section("reconstruction_eager"):
+    reconstruction = compute_reconstruction(
+        jnp.array(data_vector), jnp.array(curvature_matrix)
+    )
+    block(reconstruction)
+
+_, reconstruction = jit_profile(
+    compute_reconstruction,
+    "reconstruction_jit",
+    jnp.array(data_vector),
+    jnp.array(curvature_matrix),
+)
+likelihood_steps.append(("Reconstruction (NNLS)", timer.records[-1][1] / 10))
+
+print(f"  reconstruction shape: {reconstruction.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 7: Map reconstruction back to image plane
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 7: Mapped reconstructed image ---")
+
+
+def compute_mapped_recon(blurred_mapping_matrix, reconstruction):
+    return al.util.inversion.mapped_reconstructed_data_via_mapping_matrix_from(
+        mapping_matrix=blurred_mapping_matrix,
+        reconstruction=reconstruction,
+        xp=jnp,
+    )
+
+
+with timer.section("mapped_recon_eager"):
+    mapped_recon = compute_mapped_recon(bmm_jnp, jnp.array(reconstruction))
+    block(mapped_recon)
+
+_, mapped_recon = jit_profile(
+    compute_mapped_recon, "mapped_recon_jit", bmm_jnp, jnp.array(reconstruction)
+)
+likelihood_steps.append(("Mapped reconstructed image", timer.records[-1][1] / 10))
+
+print(f"  mapped_reconstructed_image shape: {mapped_recon.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 8: Chi-squared and log likelihood
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 8: Chi-squared & log likelihood ---")
+
+
+def compute_log_likelihood(data, noise_map, mapped_recon):
+    residual = data - mapped_recon
+    chi_squared = jnp.sum((residual / noise_map) ** 2)
+    noise_norm = jnp.sum(jnp.log(2 * jnp.pi * noise_map**2))
+    return -0.5 * (chi_squared + noise_norm)
+
+
+mapped_recon_jnp = jnp.array(mapped_recon)
+
+with timer.section("log_likelihood_eager"):
+    log_like = compute_log_likelihood(data_array, noise_jnp, mapped_recon_jnp)
+    block(log_like)
+
+_, log_like = jit_profile(
+    compute_log_likelihood,
+    "log_likelihood_jit",
+    data_array,
+    noise_jnp,
+    mapped_recon_jnp,
+)
+likelihood_steps.append(("Chi-squared & log likelihood", timer.records[-1][1] / 10))
+
+print(f"  log_likelihood = {log_like}")
+
+# Assert step-by-step result matches FitImaging.log_likelihood
+# (log_likelihood = -0.5 * (chi_squared + noise_norm), same formula as compute_log_likelihood)
+np.testing.assert_allclose(
+    float(log_like),
+    float(log_likelihood_ref),
+    rtol=1e-4,
+    err_msg="Step-by-step log_likelihood does not match FitImaging.log_likelihood",
+)
+print("  Assertion PASSED: step-by-step matches FitImaging.log_likelihood")
+
+# ===================================================================
+# PART C — Full-pipeline JIT for comparison
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("FULL-PIPELINE JIT (for comparison)")
+print("=" * 70)
+
+# Build the analysis with ``use_jax=True`` so its ``log_likelihood_function``
+# threads ``xp=jnp`` through every internal call (border relocation, profile
+# evaluation, inversion, etc.). This is the same wiring that ``Fitness.call``
+# uses in production — we just feed it our pytree-native instance directly
+# instead of going through ``model.instance_from_vector(parameters, xp=jnp)``.
+analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)
+
+
+def full_pipeline_from_params(params_tree):
+    """Full likelihood from a pytree-shaped ``ModelInstance``.
+
+    No flat-vector unpacking inside the trace — the instance crosses the JIT
+    boundary directly, with constants (redshifts, etc.) kept static via the
+    ``aux_data`` partition set up by ``autofit.jax.register_model``.
+    """
+    return analysis.log_likelihood_function(instance=params_tree)
+
+
+_, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree)
+full_pipeline_per_call = timer.records[-1][1] / 10
+
+print(f"  full log_likelihood = {full_result}")
+
+# ===================================================================
+# PART D — vmap + correctness
+# ===================================================================
+
+print("\n--- vmap batched evaluation ---")
+
+batch_size = 3
+
+# Build the batched pytree: every leaf gets a fresh leading batch axis. No
+# flat-vector reshaping required — JAX walks the pytree via the registration
+# we added in PART A.
+parameters = jax.tree_util.tree_map(
+    lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)),
+    params_tree,
+)
+
+vmapped_full = jax.jit(jax.vmap(full_pipeline_from_params))
+
+with timer.section("vmap_first_call"):
+    result_vmap = vmapped_full(parameters)
+    block(result_vmap)
+
+n_vmap_repeats = 10
+with timer.section(f"vmap_steady_x{n_vmap_repeats}"):
+    for _ in range(n_vmap_repeats):
+        result_vmap = vmapped_full(parameters)
+        block(result_vmap)
+
+vmap_batch_time = timer.records[-1][1] / n_vmap_repeats
+vmap_per_call = vmap_batch_time / batch_size
+vmap_speedup = full_pipeline_per_call / vmap_per_call
+
+print(f"  batch results = {result_vmap}")
+print(f"  vmap batch of {batch_size}:   {vmap_batch_time:.6f} s")
+print(f"  vmap per call:         {vmap_per_call:.6f} s")
+print(f"  single JIT per call:   {full_pipeline_per_call:.6f} s")
+print(f"  vmap speedup:          {vmap_speedup:.1f}x faster per likelihood")
+
+np.testing.assert_allclose(
+    np.array(result_vmap),
+    float(full_result),
+    rtol=1e-4,
+    err_msg="mge: JAX vmap likelihood mismatch",
+)
+print("  Correctness check PASSED")
+
+# ===================================================================
+# PART E — Static memory analysis
+# ===================================================================
+
+print("\n--- Static memory analysis ---")
+
+lowered_batched = vmapped_full.lower(parameters)
+compiled_batched = lowered_batched.compile()
+
+memory_analysis = compiled_batched.memory_analysis()
+print(f"  Output size:  {memory_analysis.output_size_in_bytes / 1024**2:.3f} MB")
+print(f"  Temp size:    {memory_analysis.temp_size_in_bytes / 1024**2:.3f} MB")
+print(
+    f"  Total:        "
+    f"{(memory_analysis.output_size_in_bytes + memory_analysis.temp_size_in_bytes) / 1024**2:.3f} MB"
+)
+
+
+# ===================================================================
+# JAX Likelihood Function Summary
+# ===================================================================
+
+import json
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+al_version = al.__version__
+
+print("\n" + "=" * 70)
+print(f"JAX LIKELIHOOD FUNCTION SUMMARY — {instrument.upper()} — v{al_version}")
+print("=" * 70)
+print(f"  Instrument:            {instrument}")
+print(f"  Pixel scale:           {pixel_scale} arcsec/pixel")
+print(f"  Mask radius:           {mask_radius} arcsec")
+print(f"  Image pixels (masked): {n_image_pixels}")
+print(f"  Over-sampled pixels:   {n_over_sampled_pixels}")
+print(f"  Linear Gaussians:      {n_linear_gaussians}")
+print("-" * 70)
+
+max_label = max(len(label) for label, _ in likelihood_steps)
+step_total = 0.0
+for i, (label, per_call) in enumerate(likelihood_steps, 1):
+    print(f"  {i:>2}. {label:<{max_label}}  {per_call:>12.6f} s")
+    step_total += per_call
+
+print("-" * 70)
+print(f"      {'TOTAL (step-by-step)':<{max_label}}  {step_total:>12.6f} s")
+print(
+    f"      {'Full pipeline (single JIT)':<{max_label}}  {full_pipeline_per_call:>12.6f} s"
+)
+print(
+    f"      {f'vmap batch={batch_size} (per call)':<{max_label}}  {vmap_per_call:>12.6f} s"
+)
+print(f"      {f'vmap speedup vs single JIT':<{max_label}}  {vmap_speedup:>11.1f}x")
+print("=" * 70)
+
+# --- Save results dictionary ---
+
+likelihood_summary = {
+    "autolens_version": al_version,
+    "instrument": instrument,
+    "configuration": {
+        "pixel_scale_arcsec": pixel_scale,
+        "mask_radius_arcsec": mask_radius,
+        "image_pixels_masked": int(n_image_pixels),
+        "over_sampled_pixels": int(n_over_sampled_pixels),
+        "linear_gaussians": int(n_linear_gaussians),
+    },
+    "steps": {label: per_call for label, per_call in likelihood_steps},
+    "total_step_by_step": step_total,
+    "full_pipeline_single_jit": full_pipeline_per_call,
+    "vmap": {
+        "batch_size": batch_size,
+        "batch_time": vmap_batch_time,
+        "per_call": vmap_per_call,
+        "speedup_vs_single_jit": round(vmap_speedup, 1),
+    },
+}
+
+results_dir = _workspace_root / "results" / "likelihood" / "imaging"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+dict_path = results_dir / f"mge_likelihood_summary_{instrument}_v{al_version}.json"
+dict_path.write_text(json.dumps(likelihood_summary, indent=2))
+print(f"\n  Results dict saved to: {dict_path}")
+
+# --- Save bar chart ---
+
+labels = [label for label, _ in likelihood_steps]
+times = [per_call for _, per_call in likelihood_steps]
+
+fig, ax = plt.subplots(figsize=(10, 5))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color="#4C72B0", edgecolor="white", height=0.6)
+
+for bar, t in zip(bars, times):
+    ax.text(
+        bar.get_width() + max(times) * 0.01,
+        bar.get_y() + bar.get_height() / 2,
+        f"{t:.6f} s",
+        va="center",
+        fontsize=9,
+    )
+
+ax.axvline(
+    full_pipeline_per_call,
+    color="#C44E52",
+    linestyle="--",
+    linewidth=1.5,
+    label=f"Full pipeline (single JIT): {full_pipeline_per_call:.6f} s",
+)
+ax.axvline(
+    vmap_per_call,
+    color="#55A868",
+    linestyle="--",
+    linewidth=1.5,
+    label=f"vmap batch={batch_size} per call: {vmap_per_call:.6f} s ({vmap_speedup:.1f}x faster)",
+)
+
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=10)
+ax.invert_yaxis()
+ax.set_xlabel("Time per call (s)", fontsize=11)
+fig.suptitle(
+    f"MGE Imaging Likelihood — {instrument.upper()}",
+    fontsize=12,
+    fontweight="bold",
+)
+ax.set_title(
+    f'AutoLens v{al_version}  |  {pixel_scale}"/px  |  {n_image_pixels} pixels  |  '
+    f"{n_over_sampled_pixels} over-sampled  |  {n_linear_gaussians} Gaussians  |  "
+    f"total: {step_total:.6f} s",
+    fontsize=9,
+)
+ax.legend(loc="lower right", fontsize=9)
+ax.margins(x=0.15)
+fig.tight_layout()
+
+chart_path = results_dir / f"mge_likelihood_summary_{instrument}_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to:    {chart_path}")
+
+
+# ===================================================================
+# Regression assertion — realistic-scale deterministic likelihood
+# ===================================================================
+#
+# Simulator truth parameters (mass + shear fixed; MGE bulges free around
+# default centre/ell_comps priors) put the evaluation point at the
+# physically-meaningful truth operating point. Eager, JIT, and vmap all
+# agree to ~1e-11 precision.
+EXPECTED_LOG_LIKELIHOOD_HST = 27379.38890685539
+
+np.testing.assert_allclose(
+    log_likelihood_ref,
+    EXPECTED_LOG_LIKELIHOOD_HST,
+    rtol=1e-4,
+    err_msg=(
+        f"imaging/mge[{instrument}]: regression — eager log_likelihood drifted "
+        f"(got {log_likelihood_ref}, expected {EXPECTED_LOG_LIKELIHOOD_HST})"
+    ),
+)
+print(
+    f"  Eager regression assertion PASSED: log_likelihood matches "
+    f"{EXPECTED_LOG_LIKELIHOOD_HST:.6f}"
+)
+np.testing.assert_allclose(
+    float(full_result),
+    EXPECTED_LOG_LIKELIHOOD_HST,
+    rtol=1e-4,
+    err_msg=f"imaging/mge[{instrument}]: regression — full log_likelihood drifted",
+)
+np.testing.assert_allclose(
+    np.array(result_vmap),
+    EXPECTED_LOG_LIKELIHOOD_HST,
+    rtol=1e-4,
+    err_msg=f"imaging/mge[{instrument}]: regression — vmap log_likelihood drifted",
+)
+print(
+    f"  Regression assertion PASSED: log_likelihood matches {EXPECTED_LOG_LIKELIHOOD_HST:.6f}"
+)
diff --git a/likelihood/imaging/pixelization.py b/likelihood/imaging/pixelization.py
new file mode 100644
index 0000000..8ecf0c4
--- /dev/null
+++ b/likelihood/imaging/pixelization.py
@@ -0,0 +1,1131 @@
+"""
+JAX Profiling: Pixelization Imaging Likelihood (Step-by-Step)
+=============================================================
+
+Profiles each step of the JAX likelihood function for an imaging dataset where
+the source galaxy is reconstructed using a rectangular pixelization with
+constant regularization.
+
+Rather than timing the whole likelihood as a single JIT-compiled block (which
+hides internal bottlenecks), this script JIT-compiles and times each step of
+the pipeline individually:
+
+1. Ray-trace grids through the lens
+2. Blurred image of lens light (non-linear profiles)
+3. Profile-subtracted image (lens light subtraction)
+4. Border relocation of traced grid
+5. Overlay grid (source pixel centres)
+6. Interpolation weights and mapper construction
+7. Mapping matrix
+8. Blurred mapping matrix (PSF convolution)
+9. Data vector (D)
+10. Curvature matrix (F)
+11. Regularization matrix (H)
+12. Regularized reconstruction: s = (F + H)^{-1} D
+13. Map reconstruction to image + log evidence
+
+Caveat: XLA may fuse operations differently when compiled as one program vs
+separate pieces, so per-step timings are approximate. They are still useful
+for identifying which step dominates.
+
+All JAX timings use `block_until_ready()` to force synchronous measurement.
+"""
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import time
+import subprocess
+import sys
+from pathlib import Path
+from contextlib import contextmanager
+
+import autofit as af
+import autolens as al
+import autoarray as aa
+from autofit.jax import register_model as _register_model_pytrees
+
+# ---------------------------------------------------------------------------
+# Instrument configuration
+# ---------------------------------------------------------------------------
+
+INSTRUMENTS = {
+    "euclid": {"pixel_scale": 0.1},
+    "hst": {"pixel_scale": 0.05},
+    "jwst": {"pixel_scale": 0.03},
+    "ao": {"pixel_scale": 0.01},
+}
+
+instrument = "hst"  # <-- change this to profile a different instrument
+
+
+# ---------------------------------------------------------------------------
+# Profiling helpers
+# ---------------------------------------------------------------------------
+
+
+class Timer:
+    """Accumulates named timing measurements and prints a summary."""
+
+    def __init__(self):
+        self.records: list[tuple[str, float]] = []
+
+    @contextmanager
+    def section(self, label: str):
+        """Context manager that records wall-clock time for *label*."""
+        start = time.perf_counter()
+        yield
+        elapsed = time.perf_counter() - start
+        self.records.append((label, elapsed))
+        print(f"  [{label}] {elapsed:.4f} s")
+
+    def summary(self):
+        print("\n" + "=" * 70)
+        print("PROFILING SUMMARY")
+        print("=" * 70)
+        max_label = max(len(r[0]) for r in self.records)
+        total = 0.0
+        for label, elapsed in self.records:
+            print(f"  {label:<{max_label}}  {elapsed:>10.4f} s")
+            total += elapsed
+        print("-" * 70)
+        print(f"  {'TOTAL':<{max_label}}  {total:>10.4f} s")
+        print("=" * 70)
+
+
+def block(x):
+    """Call block_until_ready if available (JAX arrays)."""
+    if hasattr(x, "block_until_ready"):
+        x.block_until_ready()
+    return x
+
+
+def jit_profile(func, label, *args, n_repeats=10):
+    """JIT-compile *func*, time first call and steady-state average.
+
+    Returns the compiled function and its result.
+    """
+    jitted = jax.jit(func)
+
+    with timer.section(f"{label}_lower"):
+        lowered = jitted.lower(*args)
+
+    with timer.section(f"{label}_compile"):
+        compiled = lowered.compile()
+
+    with timer.section(f"{label}_first_call"):
+        result = compiled(*args)
+        block(result)
+
+    with timer.section(f"{label}_steady_x{n_repeats}"):
+        for _ in range(n_repeats):
+            result = compiled(*args)
+            block(result)
+
+    per_call = timer.records[-1][1] / n_repeats
+    print(f"    -> per-call avg: {per_call:.6f} s")
+    return compiled, result
+
+
+timer = Timer()
+likelihood_steps = []  # (label, per_call_seconds) for the final summary
+
+# ===================================================================
+# PART A — Setup (not JIT-compiled)
+# ===================================================================
+
+# ---------------------------------------------------------------------------
+# 1. Dataset
+# ---------------------------------------------------------------------------
+
+print(f"\n--- Dataset loading & masking [{instrument}] ---")
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[1]
+pixel_scale = INSTRUMENTS[instrument]["pixel_scale"]
+dataset_path = Path("dataset") / "imaging" / instrument
+
+if al.util.dataset.should_simulate(str(dataset_path)):
+    raise FileNotFoundError(
+        f"Input dataset missing at '{dataset_path}'. The autolens_profiling "
+        f"repo mirrors only the curated datasets needed for default smoke "
+        f"runs. To regenerate or extend datasets, use the source-of-truth "
+        f"scripts under autolens_workspace_developer/jax_profiling/dataset_setup/, "
+        f"then copy the result into autolens_profiling/dataset/."
+    )
+
+with timer.section("dataset_load"):
+    dataset = al.Imaging.from_fits(
+        data_path=dataset_path / "data.fits",
+        psf_path=dataset_path / "psf.fits",
+        noise_map_path=dataset_path / "noise_map.fits",
+        pixel_scales=pixel_scale,
+    )
+
+with timer.section("mask_and_oversample"):
+    mask_radius = 3.5
+
+    mask = al.Mask2D.circular(
+        shape_native=dataset.shape_native,
+        pixel_scales=dataset.pixel_scales,
+        radius=mask_radius,
+    )
+
+    dataset = dataset.apply_mask(mask=mask)
+    dataset = dataset.apply_over_sampling(
+        over_sample_size_lp=4,
+        over_sample_size_pixelization=1,
+    )
+
+    over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(
+        grid=dataset.grid,
+        sub_size_list=[4, 2, 1],
+        radial_list=[0.3, 0.6],
+        centre_list=[(0.0, 0.0)],
+    )
+
+    dataset = dataset.apply_over_sampling(
+        over_sample_size_lp=over_sample_size,
+        over_sample_size_pixelization=1,
+    )
+
+# ---------------------------------------------------------------------------
+# 2. Model construction
+# ---------------------------------------------------------------------------
+
+print("\n--- Model construction ---")
+
+mesh_pixels_yx = 35  # 35x35 = 1225, science fiducial near 1250 source pixels
+mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)
+
+with timer.section("model_build"):
+    # GaussianPrior(mean=truth, sigma=small) centres prior-median at the
+    # simulator truth while keeping params free so gradient diagnostics
+    # have dimensionality.
+    # Lens light: MGE-60 (full production-fiducial) — replaces single Sersic.
+    # The 60 linear Gaussians enter the inversion's mapping matrix
+    # alongside the source-pixel columns.
+    lens_bulge = al.model_util.mge_model_from(
+        mask_radius=mask_radius,
+        total_gaussians=60,
+        centre_prior_is_uniform=True,
+    )
+
+    mass = af.Model(al.mp.Isothermal)
+    mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.005)
+    mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.005)
+    mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05)
+    _lens_mass_ell = al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0)
+    mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=_lens_mass_ell[0], sigma=0.01)
+    mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=_lens_mass_ell[1], sigma=0.01)
+
+    shear = af.Model(al.mp.ExternalShear)
+    shear.gamma_1 = af.GaussianPrior(mean=0.05, sigma=0.005)
+    shear.gamma_2 = af.GaussianPrior(mean=0.05, sigma=0.005)
+
+    lens = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear)
+
+    pixelization = al.Pixelization(
+        mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),
+        regularization=al.reg.Constant(coefficient=1.0),
+    )
+
+    source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)
+
+    model = af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+print(f"  Total free parameters: {model.total_free_parameters}")
+print(f"  Mesh shape: {mesh_shape}")
+print(f"  Source pixels: {mesh_pixels_yx * mesh_pixels_yx}")
+
+# ---------------------------------------------------------------------------
+# 3. Instantiate concrete objects from prior medians
+# ---------------------------------------------------------------------------
+
+print("\n--- Instantiate concrete model ---")
+
+with timer.section("instance_from_vector"):
+    param_vector = model.physical_values_from_prior_medians
+    instance = model.instance_from_vector(vector=param_vector)
+
+with timer.section("register_pytrees"):
+    _register_model_pytrees(model)
+
+params_tree = jax.tree_util.tree_map(jnp.asarray, instance)
+tracer = al.Tracer(galaxies=list(instance.galaxies))
+
+print(f"  Tracer planes: {tracer.total_planes}")
+
+# ---------------------------------------------------------------------------
+# Key configuration that dictates run time
+# ---------------------------------------------------------------------------
+
+n_image_pixels = dataset.data.shape[0]
+n_over_sampled_pixels = dataset.grids.lp.over_sampled.shape[0]
+n_source_pixels = mesh_pixels_yx * mesh_pixels_yx
+
+print("\n--- Configuration (determines run time) ---")
+print(f"  Instrument:              {instrument}")
+print(f"  Pixel scale:             {pixel_scale} arcsec/pixel")
+print(f"  Mask radius:             {mask_radius} arcsec")
+print(f"  Image pixels (masked):   {n_image_pixels}")
+print(f"  Over-sampled pixels:     {n_over_sampled_pixels}")
+print(f"  Mesh shape:              {mesh_shape}")
+print(f"  Source pixels:           {n_source_pixels}")
+
+# ---------------------------------------------------------------------------
+# 4. Full-pipeline reference (FitImaging) — eager baseline
+# ---------------------------------------------------------------------------
+
+print("\n--- Full FitImaging (eager baseline) ---")
+
+with timer.section("fit_imaging_eager"):
+    fit = al.FitImaging(
+        dataset=dataset,
+        tracer=tracer,
+        settings=al.Settings(use_border_relocator=True),
+        xp=np,
+    )
+    log_evidence_ref = fit.figure_of_merit
+    log_likelihood_ref = fit.log_likelihood
+
+print(f"  figure_of_merit (log_evidence) = {log_evidence_ref}")
+print(f"  log_likelihood                 = {log_likelihood_ref}")
+
+
+# ===================================================================
+# PART B — Per-step JIT profiling
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("PER-STEP JIT PROFILING")
+print("=" * 70)
+
+# Extract raw arrays from autoarray types via .array so they can cross
+# JIT boundaries.  See CLAUDE.md for rationale.
+
+grid_lp_raw = jnp.array(dataset.grids.lp.array)
+grid_pix_raw = jnp.array(dataset.grids.pixelization.array)
+grid_blurring_raw = jnp.array(dataset.grids.blurring.array)
+data_array = jnp.array(dataset.data.array)
+noise_map_array = jnp.array(dataset.noise_map.array)
+
+# Keep autoarray objects for eager calls that need them.
+grid_lp = dataset.grids.lp
+grid_blurring = dataset.grids.blurring
+
+# ---------------------------------------------------------------------------
+# Step 1: Ray-trace grids
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 1: Ray-trace grids ---")
+
+with timer.section("ray_trace_eager"):
+    traced_grids = tracer.traced_grid_2d_list_from(
+        grid=dataset.grids.pixelization, xp=jnp
+    )
+    for tg in traced_grids:
+        block(tg)
+
+print(f"  Number of planes traced: {len(traced_grids)}")
+
+
+def ray_trace_raw(grid_raw):
+    """Wraps ray-tracing so inputs/outputs are raw arrays."""
+    grid = aa.Grid2DIrregular(values=grid_raw, xp=jnp)
+    traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp)
+    return jnp.stack([tg.array for tg in traced])
+
+
+_, traced_grids_raw = jit_profile(ray_trace_raw, "ray_trace_jit", grid_pix_raw)
+likelihood_steps.append(("Ray-trace grids", timer.records[-1][1] / 10))
+
+print(f"  traced_grids shape: {traced_grids_raw.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 2: Blurred image of non-linear light profiles (lens light)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 2: Blurred image (lens light profiles) ---")
+
+
+# Sub-step 2a: Compute raw lens light images (JIT-profiled)
+def lens_image_raw(grid_raw, blurring_grid_raw):
+    """Compute lens light images on masked + blurring grids (no PSF)."""
+    grid = aa.Grid2DIrregular(values=grid_raw, xp=jnp)
+    blurring_grid = aa.Grid2DIrregular(values=blurring_grid_raw, xp=jnp)
+    image = tracer.image_2d_from(grid=grid, xp=jnp)
+    blurring_image = tracer.image_2d_from(grid=blurring_grid, xp=jnp)
+    return image.array, blurring_image.array
+
+
+with timer.section("lens_image_eager"):
+    img_eager, blur_img_eager = lens_image_raw(grid_lp_raw, grid_blurring_raw)
+    block(img_eager)
+    block(blur_img_eager)
+
+_, (img_jit, blur_img_jit) = jit_profile(
+    lens_image_raw, "lens_image_jit", grid_lp_raw, grid_blurring_raw
+)
+likelihood_steps.append(("Lens light images (pre-PSF)", timer.records[-1][1] / 10))
+
+# Sub-step 2b: PSF convolution (eager — requires autoarray mask objects)
+with timer.section("blurred_image_eager"):
+    blurred_image = tracer.blurred_image_2d_from(
+        grid=grid_lp,
+        psf=dataset.psf,
+        blurring_grid=grid_blurring,
+        xp=jnp,
+    )
+    block(blurred_image)
+
+print(f"  blurred_image shape: {blurred_image.array.shape}")
+
+
+def blurred_image_from_params(params_tree):
+    """Compute blurred image directly from a pytree ModelInstance — fully JIT-traceable."""
+    t = al.Tracer(galaxies=list(params_tree.galaxies))
+    result = t.blurred_image_2d_from(
+        grid=grid_lp,
+        psf=dataset.psf,
+        blurring_grid=grid_blurring,
+        xp=jnp,
+    )
+    return result.array
+
+
+_, blurred_img_jit = jit_profile(
+    blurred_image_from_params, "blurred_image_jit", params_tree
+)
+likelihood_steps.append(("Blurred image (PSF convolution)", timer.records[-1][1] / 10))
+
+# ---------------------------------------------------------------------------
+# Step 3: Profile-subtracted image (lens light subtraction)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 3: Profile-subtracted image ---")
+
+
+def profile_subtract(data, blurred_image):
+    return data - blurred_image
+
+
+with timer.section("profile_subtract_eager"):
+    blurred_img_jnp = jnp.array(blurred_image.array)
+    profile_subtracted = profile_subtract(data_array, blurred_img_jnp)
+    block(profile_subtracted)
+
+_, profile_subtracted = jit_profile(
+    profile_subtract, "profile_subtract_jit", data_array, blurred_img_jnp
+)
+likelihood_steps.append(("Profile-subtracted image", timer.records[-1][1] / 10))
+
+print(f"  profile_subtracted shape: {profile_subtracted.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 4: Border relocation of traced grid
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 4: Border relocation ---")
+
+from autoarray.inversion.mesh.border_relocator import BorderRelocator
+
+with timer.section("border_relocator_setup"):
+    border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1)
+
+# The source plane grid is the last entry (index -1) of the traced grids list.
+traced_source_grid = tracer.traced_grid_2d_list_from(
+    grid=dataset.grids.pixelization, xp=jnp
+)[-1]
+
+with timer.section("border_relocation_eager"):
+    relocated_grid = border_relocator.relocated_grid_from(grid=traced_source_grid)
+    block(relocated_grid)
+
+print(f"  relocated_grid shape: {relocated_grid.array.shape}")
+
+# For JIT profiling, extract the relocation logic as a raw-array function
+relocated_grid_raw = jnp.array(relocated_grid.array)
+
+# ---------------------------------------------------------------------------
+# Step 5: Overlay grid (source pixel centres)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 5: Overlay grid (source pixel centres) ---")
+
+from autoarray.inversion.mesh.mesh.rectangular_adapt_density import overlay_grid_from
+
+with timer.section("overlay_grid_eager"):
+    mesh_grid = overlay_grid_from(
+        shape_native=mesh_shape,
+        grid=al.Grid2DIrregular(relocated_grid),
+        xp=jnp,
+    )
+    block(mesh_grid)
+
+
+def overlay_grid_raw_fn(relocated_grid_raw):
+    grid = al.Grid2DIrregular(values=relocated_grid_raw, xp=jnp)
+    return overlay_grid_from(shape_native=mesh_shape, grid=grid, xp=jnp)
+
+
+_, mesh_grid_raw = jit_profile(
+    overlay_grid_raw_fn, "overlay_grid_jit", relocated_grid_raw
+)
+likelihood_steps.append(
+    ("Overlay grid (source pixel centres)", timer.records[-1][1] / 10)
+)
+
+print(f"  mesh_grid shape: {mesh_grid_raw.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 6: Interpolation + Mapper construction
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 6: Interpolation + Mapper ---")
+
+pixelization_obj = instance.galaxies.source.pixelization
+
+with timer.section("interpolation_and_mapper"):
+    interpolator = pixelization_obj.mesh.interpolator_from(
+        source_plane_data_grid=relocated_grid,
+        source_plane_mesh_grid=al.Grid2DIrregular(mesh_grid),
+    )
+    mapper = al.Mapper(interpolator=interpolator, xp=jnp)
+
+print(f"  mapper.pixels (source): {mapper.pixels}")
+print(f"  pix_indexes shape: {mapper.pix_indexes_for_sub_slim_index.shape}")
+
+# ---------------------------------------------------------------------------
+# Steps 7-13: Extract matrices from FitImaging inversion for consistency
+# ---------------------------------------------------------------------------
+# The FitImaging pipeline handles edge pixel zeroing, curvature diagonal
+# adjustments, and settings that are difficult to replicate manually.
+# We extract the correct matrices from fit.inversion so the step-by-step
+# matches the reference, then JIT-profile the linear algebra operations.
+
+print("\n--- Extracting inversion matrices from FitImaging ---")
+
+inversion = fit.inversion
+
+with timer.section("extract_inversion_matrices"):
+    # The operated_mapping_matrix is the blurred mapping matrix (post-PSF)
+    bmm_ref = jnp.array(inversion.operated_mapping_matrix)
+    mapping_matrix_ref = jnp.array(inversion.mapping_matrix)
+
+    # Extract mapper from the inversion's linear object list
+    inv_mapper = inversion.cls_list_from(cls=al.Mapper)[0]
+    neighbors = inv_mapper.neighbors
+    neighbors_array = jnp.array(np.asarray(neighbors))
+    neighbors_sizes = jnp.array(neighbors.sizes)
+
+    reg_coefficient = pixelization_obj.regularization.coefficient
+
+print(f"  operated_mapping_matrix shape: {bmm_ref.shape}")
+print(f"  mapping_matrix shape: {mapping_matrix_ref.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 7: Mapping matrix
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 7: Mapping matrix ---")
+
+with timer.section("mapping_matrix"):
+    mapping_matrix = inv_mapper.mapping_matrix
+
+print(f"  mapping_matrix shape: {mapping_matrix.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 8: Blurred mapping matrix (PSF convolution)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 8: Blurred mapping matrix ---")
+
+with timer.section("blurred_mapping_matrix"):
+    blurred_mapping_matrix = dataset.psf.convolved_mapping_matrix_from(
+        mapping_matrix=mapping_matrix,
+        mask=dataset.mask,
+        xp=jnp,
+    )
+    block(blurred_mapping_matrix)
+
+# JIT-profile the full inversion setup pipeline (steps 4-8 combined):
+# border relocation → overlay grid → interpolation → mapper → mapping matrix → PSF convolution.
+# These steps are tightly sequential; the full pipeline JIT-compiles them all together.
+
+
+def blurred_mm_from_params(params_tree):
+    """Compute blurred mapping matrix via full inversion setup from a pytree ModelInstance."""
+    t = al.Tracer(galaxies=list(params_tree.galaxies))
+    fit_jax = al.FitImaging(
+        dataset=dataset,
+        tracer=t,
+        settings=al.Settings(use_border_relocator=True),
+        xp=jnp,
+    )
+    return jnp.array(fit_jax.inversion.operated_mapping_matrix)
+
+
+_, bmm_jit = jit_profile(blurred_mm_from_params, "inversion_setup_jit", params_tree)
+likelihood_steps.append(
+    ("Inversion setup (steps 4-8 combined)", timer.records[-1][1] / 10)
+)
+
+print(f"  blurred_mapping_matrix (JIT) shape: {bmm_jit.shape}")
+
+bmm_jnp = bmm_ref  # Use the reference matrices for linear algebra steps
+print(f"  blurred_mapping_matrix shape: {blurred_mapping_matrix.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 9: Data vector (D)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 9: Data vector ---")
+
+
+def compute_data_vector(blurred_mapping_matrix, image, noise_map):
+    return al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from(
+        blurred_mapping_matrix=blurred_mapping_matrix,
+        image=image,
+        noise_map=noise_map,
+    )
+
+
+profile_sub_jnp = jnp.array(fit.profile_subtracted_image.array)
+noise_jnp = jnp.array(dataset.noise_map.array)
+
+with timer.section("data_vector_eager"):
+    data_vector = compute_data_vector(bmm_jnp, profile_sub_jnp, noise_jnp)
+    block(data_vector)
+
+_, data_vector = jit_profile(
+    compute_data_vector, "data_vector_jit", bmm_jnp, profile_sub_jnp, noise_jnp
+)
+likelihood_steps.append(("Data vector (D)", timer.records[-1][1] / 10))
+
+print(f"  data_vector shape: {data_vector.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 10: Curvature matrix (F)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 10: Curvature matrix ---")
+
+# Match the FitImaging inversion: add_to_curvature_diag=True, with settings
+no_reg_list = list(inversion.no_regularization_index_list)
+
+
+def compute_curvature_matrix(blurred_mapping_matrix, noise_map):
+    return al.util.inversion.curvature_matrix_via_mapping_matrix_from(
+        mapping_matrix=blurred_mapping_matrix,
+        noise_map=noise_map,
+        settings=fit.settings,
+        add_to_curvature_diag=True,
+        no_regularization_index_list=no_reg_list,
+        xp=jnp,
+    )
+
+
+with timer.section("curvature_matrix_eager"):
+    curvature_matrix = compute_curvature_matrix(bmm_jnp, noise_jnp)
+    block(curvature_matrix)
+
+_, curvature_matrix = jit_profile(
+    compute_curvature_matrix, "curvature_matrix_jit", bmm_jnp, noise_jnp
+)
+likelihood_steps.append(("Curvature matrix (F)", timer.records[-1][1] / 10))
+
+print(f"  curvature_matrix shape: {curvature_matrix.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 11: Regularization matrix (H)
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 11: Regularization matrix ---")
+
+
+def compute_regularization_matrix(neighbors_array, neighbors_sizes):
+    return al.util.regularization.constant_regularization_matrix_from(
+        coefficient=reg_coefficient,
+        neighbors=neighbors_array,
+        neighbors_sizes=neighbors_sizes,
+        xp=jnp,
+    )
+
+
+with timer.section("regularization_matrix_eager"):
+    regularization_matrix = compute_regularization_matrix(
+        neighbors_array, neighbors_sizes
+    )
+    block(regularization_matrix)
+
+_, regularization_matrix = jit_profile(
+    compute_regularization_matrix,
+    "regularization_matrix_jit",
+    neighbors_array,
+    neighbors_sizes,
+)
+likelihood_steps.append(("Regularization matrix (H)", timer.records[-1][1] / 10))
+
+print(f"  regularization_matrix shape: {regularization_matrix.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 12: Regularized reconstruction: s = (F + H)^{-1} D
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 12: Regularized reconstruction ---")
+
+
+def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix):
+    curvature_reg_matrix = curvature_matrix + regularization_matrix
+    return al.util.inversion.reconstruction_positive_only_from(
+        data_vector=data_vector,
+        curvature_reg_matrix=curvature_reg_matrix,
+        xp=jnp,
+    )
+
+
+with timer.section("reconstruction_eager"):
+    reconstruction = compute_reconstruction(
+        jnp.array(data_vector),
+        jnp.array(curvature_matrix),
+        jnp.array(regularization_matrix),
+    )
+    block(reconstruction)
+
+_, reconstruction = jit_profile(
+    compute_reconstruction,
+    "reconstruction_jit",
+    jnp.array(data_vector),
+    jnp.array(curvature_matrix),
+    jnp.array(regularization_matrix),
+)
+likelihood_steps.append(("Regularized reconstruction", timer.records[-1][1] / 10))
+
+print(f"  reconstruction shape: {reconstruction.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 13: Map reconstruction to image + log evidence
+# ---------------------------------------------------------------------------
+
+print("\n--- Step 13: Mapped reconstruction + log evidence ---")
+
+
+def compute_log_evidence(
+    data,
+    noise_map,
+    blurred_image,
+    blurred_mapping_matrix,
+    reconstruction,
+    curvature_matrix,
+    regularization_matrix,
+    mapper_indices,
+):
+    """Compute the full log evidence including all five terms:
+
+    -2 ln e = chi^2 + s^T H s + ln[det(F+H)] - ln[det(H)] + noise_norm
+
+    Matches the production formula in
+    ``autoarray/inversion/inversion/abstract.py:log_det_*`` —
+    reduces both the curvature_reg_matrix and the regularization_matrix
+    to the rows/cols indexed by ``mapper_indices`` before the log_det.
+    This drops the no-regularization rows (e.g. MGE Basis linear
+    components) which otherwise make ``det(H) = 0`` and the slogdet
+    return -inf, then uses Cholesky for a numerically stable log_det.
+    """
+    # Map reconstruction to image
+    mapped_recon = al.util.inversion.mapped_reconstructed_data_via_mapping_matrix_from(
+        mapping_matrix=blurred_mapping_matrix,
+        reconstruction=reconstruction,
+        xp=jnp,
+    )
+
+    # model_data = lens light + pixelized source
+    model_data = blurred_image + mapped_recon
+
+    # Chi-squared
+    residual = data - model_data
+    chi_squared = jnp.sum((residual / noise_map) ** 2)
+
+    # Regularization term: s^T H s
+    regularization_term = jnp.dot(
+        reconstruction, jnp.dot(regularization_matrix, reconstruction)
+    )
+
+    # Curvature + regularization matrix
+    curvature_reg_matrix = curvature_matrix + regularization_matrix
+
+    # Reduce to pixelization rows/cols only (matches production
+    # ``*_matrix_reduced``): required for models with non-regularised
+    # linear components (e.g. MGE lens light).
+    creg_reduced = curvature_reg_matrix[mapper_indices][:, mapper_indices]
+    reg_reduced = regularization_matrix[mapper_indices][:, mapper_indices]
+
+    log_det_curvature_reg = 2.0 * jnp.sum(
+        jnp.log(jnp.diag(jnp.linalg.cholesky(creg_reduced)))
+    )
+    log_det_regularization = 2.0 * jnp.sum(
+        jnp.log(jnp.diag(jnp.linalg.cholesky(reg_reduced)))
+    )
+
+    # Noise normalization
+    noise_normalization = jnp.sum(jnp.log(2 * jnp.pi * noise_map**2))
+
+    return -0.5 * (
+        chi_squared
+        + regularization_term
+        + log_det_curvature_reg
+        - log_det_regularization
+        + noise_normalization
+    )
+
+
+# For the JIT profiling we use the step-by-step matrices for timing.
+# For the correctness assertion we use the inversion's own matrices, because
+# cumulative floating-point differences between JIT-compiled and eager paths
+# (especially through ill-conditioned solves) can compound significantly.
+
+blurred_img_jnp = jnp.array(blurred_image.array)
+recon_jnp = jnp.array(reconstruction)
+curv_jnp = jnp.array(curvature_matrix)
+reg_jnp = jnp.array(regularization_matrix)
+mapper_indices_jnp = jnp.array(np.asarray(inversion.mapper_indices))
+
+with timer.section("log_evidence_eager"):
+    log_evidence = compute_log_evidence(
+        data_array,
+        noise_jnp,
+        blurred_img_jnp,
+        bmm_jnp,
+        recon_jnp,
+        curv_jnp,
+        reg_jnp,
+        mapper_indices_jnp,
+    )
+    block(log_evidence)
+
+_, log_evidence = jit_profile(
+    compute_log_evidence,
+    "log_evidence_jit",
+    data_array,
+    noise_jnp,
+    blurred_img_jnp,
+    bmm_jnp,
+    recon_jnp,
+    curv_jnp,
+    reg_jnp,
+    mapper_indices_jnp,
+)
+likelihood_steps.append(("Mapped recon + log evidence", timer.records[-1][1] / 10))
+
+print(f"  log_evidence (step-by-step) = {log_evidence}")
+
+# Correctness check: recompute log_evidence using the inversion's own
+# reconstruction and curvature matrix to avoid accumulated FP drift.
+inv_recon_jnp = jnp.array(inversion.reconstruction)
+inv_curv_jnp = jnp.array(inversion.curvature_matrix)
+
+log_evidence_check = compute_log_evidence(
+    data_array,
+    noise_jnp,
+    blurred_img_jnp,
+    bmm_jnp,
+    inv_recon_jnp,
+    inv_curv_jnp,
+    reg_jnp,
+    mapper_indices_jnp,
+)
+print(f"  log_evidence (inv matrices) = {log_evidence_check}")
+print(f"  log_evidence (reference)    = {log_evidence_ref}")
+
+np.testing.assert_allclose(
+    float(log_evidence_check),
+    float(log_evidence_ref),
+    rtol=1e-4,
+    err_msg="Log_evidence from inversion matrices does not match FitImaging.log_evidence",
+)
+print(
+    "  Assertion PASSED: inversion-matrix log_evidence matches FitImaging.log_evidence"
+)
+
+# ===================================================================
+# PART C — Full-pipeline JIT for comparison
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("FULL-PIPELINE JIT (for comparison)")
+print("=" * 70)
+
+analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)
+
+
+def full_pipeline_from_params(params_tree):
+    return analysis.log_likelihood_function(instance=params_tree)
+
+
+_, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree)
+full_pipeline_per_call = timer.records[-1][1] / 10
+
+print(f"  full log_likelihood = {full_result}")
+
+# ===================================================================
+# PART D — vmap + correctness
+# ===================================================================
+#
+# NOTE: vmap requires at least one JAX array leaf in the params_tree.
+# When model.total_free_parameters == 0 (all params fixed to truth), the
+# pytree has no array leaves and vmap cannot batch over it. Skip in that case.
+
+print("\n--- vmap batched evaluation ---")
+
+batch_size = 3
+vmap_batch_time = None
+vmap_per_call = None
+vmap_speedup = None
+result_vmap = None
+
+_n_leaves = len(jax.tree_util.tree_leaves(params_tree))
+if _n_leaves == 0:
+    print(
+        f"  SKIPPED: model has 0 free parameters (all fixed to truth); "
+        f"vmap requires at least one array leaf."
+    )
+else:
+    parameters = jax.tree_util.tree_map(
+        lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)),
+        params_tree,
+    )
+
+    vmapped_full = jax.jit(jax.vmap(full_pipeline_from_params))
+
+    with timer.section("vmap_first_call"):
+        result_vmap = vmapped_full(parameters)
+        block(result_vmap)
+
+    n_vmap_repeats = 10
+    with timer.section(f"vmap_steady_x{n_vmap_repeats}"):
+        for _ in range(n_vmap_repeats):
+            result_vmap = vmapped_full(parameters)
+            block(result_vmap)
+
+    vmap_batch_time = timer.records[-1][1] / n_vmap_repeats
+    vmap_per_call = vmap_batch_time / batch_size
+    vmap_speedup = full_pipeline_per_call / vmap_per_call
+
+    print(f"  batch results = {result_vmap}")
+    print(f"  vmap batch of {batch_size}:   {vmap_batch_time:.6f} s")
+    print(f"  vmap per call:         {vmap_per_call:.6f} s")
+    print(f"  single JIT per call:   {full_pipeline_per_call:.6f} s")
+    print(f"  vmap speedup:          {vmap_speedup:.1f}x faster per likelihood")
+
+    np.testing.assert_allclose(
+        np.array(result_vmap),
+        float(full_result),
+        rtol=1e-4,
+        err_msg="pixelization: JAX vmap likelihood mismatch",
+    )
+    print("  Correctness check PASSED")
+
+# ===================================================================
+# PART E — Static memory analysis
+# ===================================================================
+
+print("\n--- Static memory analysis ---")
+
+if _n_leaves == 0:
+    print("  SKIPPED: no array leaves in params_tree (all params fixed to truth).")
+    memory_analysis = None
+else:
+    lowered_batched = vmapped_full.lower(parameters)
+    compiled_batched = lowered_batched.compile()
+
+    memory_analysis = compiled_batched.memory_analysis()
+    print(f"  Output size:  {memory_analysis.output_size_in_bytes / 1024**2:.3f} MB")
+    print(f"  Temp size:    {memory_analysis.temp_size_in_bytes / 1024**2:.3f} MB")
+    print(
+        f"  Total:        "
+        f"{(memory_analysis.output_size_in_bytes + memory_analysis.temp_size_in_bytes) / 1024**2:.3f} MB"
+    )
+
+
+# ===================================================================
+# JAX Likelihood Function Summary
+# ===================================================================
+
+import json
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+al_version = al.__version__
+
+print("\n" + "=" * 70)
+print(f"JAX LIKELIHOOD FUNCTION SUMMARY — {instrument.upper()} — v{al_version}")
+print("=" * 70)
+print(f"  Instrument:            {instrument}")
+print(f"  Pixel scale:           {pixel_scale} arcsec/pixel")
+print(f"  Mask radius:           {mask_radius} arcsec")
+print(f"  Image pixels (masked): {n_image_pixels}")
+print(f"  Over-sampled pixels:   {n_over_sampled_pixels}")
+print(f"  Mesh shape:            {mesh_shape}")
+print(f"  Source pixels:         {n_source_pixels}")
+print("-" * 70)
+
+max_label = max(len(label) for label, _ in likelihood_steps)
+step_total = 0.0
+for i, (label, per_call) in enumerate(likelihood_steps, 1):
+    print(f"  {i:>2}. {label:<{max_label}}  {per_call:>12.6f} s")
+    step_total += per_call
+
+print("-" * 70)
+print(f"      {'TOTAL (step-by-step)':<{max_label}}  {step_total:>12.6f} s")
+print(
+    f"      {'Full pipeline (single JIT)':<{max_label}}  {full_pipeline_per_call:>12.6f} s"
+)
+if vmap_per_call is not None:
+    print(
+        f"      {f'vmap batch={batch_size} (per call)':<{max_label}}  {vmap_per_call:>12.6f} s"
+    )
+    print(f"      {f'vmap speedup vs single JIT':<{max_label}}  {vmap_speedup:>11.1f}x")
+else:
+    print(f"      {'vmap':<{max_label}}  {'SKIPPED (0 free params)':>12}")
+print("=" * 70)
+
+# --- Save results dictionary ---
+
+likelihood_summary = {
+    "autolens_version": al_version,
+    "instrument": instrument,
+    "configuration": {
+        "pixel_scale_arcsec": pixel_scale,
+        "mask_radius_arcsec": mask_radius,
+        "image_pixels_masked": int(n_image_pixels),
+        "over_sampled_pixels": int(n_over_sampled_pixels),
+        "mesh_shape": list(mesh_shape),
+        "source_pixels": int(n_source_pixels),
+    },
+    "steps": {label: per_call for label, per_call in likelihood_steps},
+    "total_step_by_step": step_total,
+    "full_pipeline_single_jit": full_pipeline_per_call,
+    "vmap": (
+        "SKIPPED — model has 0 free parameters (all fixed to truth)"
+        if vmap_per_call is None
+        else {
+            "batch_size": batch_size,
+            "batch_time": vmap_batch_time,
+            "per_call": vmap_per_call,
+            "speedup_vs_single_jit": round(vmap_speedup, 1),
+        }
+    ),
+}
+
+results_dir = _workspace_root / "results" / "likelihood" / "imaging"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+dict_path = (
+    results_dir / f"pixelization_likelihood_summary_{instrument}_v{al_version}.json"
+)
+dict_path.write_text(json.dumps(likelihood_summary, indent=2))
+print(f"\n  Results dict saved to: {dict_path}")
+
+# --- Save bar chart ---
+
+labels = [label for label, _ in likelihood_steps]
+times = [per_call for _, per_call in likelihood_steps]
+
+fig, ax = plt.subplots(figsize=(10, 6))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color="#4C72B0", edgecolor="white", height=0.6)
+
+for bar, t in zip(bars, times):
+    ax.text(
+        bar.get_width() + max(times) * 0.01,
+        bar.get_y() + bar.get_height() / 2,
+        f"{t:.6f} s",
+        va="center",
+        fontsize=9,
+    )
+
+ax.axvline(
+    full_pipeline_per_call,
+    color="#C44E52",
+    linestyle="--",
+    linewidth=1.5,
+    label=f"Full pipeline (single JIT): {full_pipeline_per_call:.6f} s",
+)
+if vmap_per_call is not None:
+    ax.axvline(
+        vmap_per_call,
+        color="#55A868",
+        linestyle="--",
+        linewidth=1.5,
+        label=f"vmap batch={batch_size} per call: {vmap_per_call:.6f} s ({vmap_speedup:.1f}x faster)",
+    )
+
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=10)
+ax.invert_yaxis()
+ax.set_xlabel("Time per call (s)", fontsize=11)
+fig.suptitle(
+    f"Pixelization Imaging Likelihood — {instrument.upper()}",
+    fontsize=12,
+    fontweight="bold",
+)
+ax.set_title(
+    f'AutoLens v{al_version}  |  {pixel_scale}"/px  |  {n_image_pixels} pixels  |  '
+    f"{n_over_sampled_pixels} over-sampled  |  {mesh_shape[0]}x{mesh_shape[1]} mesh  |  "
+    f"total: {step_total:.6f} s",
+    fontsize=9,
+)
+ax.legend(loc="lower right", fontsize=9)
+ax.margins(x=0.15)
+fig.tight_layout()
+
+chart_path = (
+    results_dir / f"pixelization_likelihood_summary_{instrument}_v{al_version}.png"
+)
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to:    {chart_path}")
+
+
+# ===================================================================
+# Regression assertion — realistic-scale deterministic log-evidence
+# ===================================================================
+#
+# RectangularAdaptDensity at prior medians is deterministic across the
+# eager / full-JIT / vmap paths to within rtol=1e-4 — the constant below
+# is the value those three paths agree on.
+EXPECTED_LOG_EVIDENCE_HST = (
+    24746.105672366088  # 35x35 = 1225 source pixels, MGE-60 lens light
+)
+
+np.testing.assert_allclose(
+    log_evidence_ref,
+    EXPECTED_LOG_EVIDENCE_HST,
+    rtol=1e-4,
+    err_msg=(
+        f"imaging/pixelization[{instrument}]: regression — eager log_evidence drifted "
+        f"(got {log_evidence_ref}, expected {EXPECTED_LOG_EVIDENCE_HST})"
+    ),
+)
+print(
+    f"  Eager regression assertion PASSED: log_evidence matches "
+    f"{EXPECTED_LOG_EVIDENCE_HST:.6f}"
+)
+np.testing.assert_allclose(
+    float(full_result),
+    EXPECTED_LOG_EVIDENCE_HST,
+    rtol=1e-4,
+    err_msg=f"imaging/pixelization[{instrument}]: regression — full log_evidence drifted",
+)
+np.testing.assert_allclose(
+    np.array(result_vmap),
+    EXPECTED_LOG_EVIDENCE_HST,
+    rtol=1e-4,
+    err_msg=f"imaging/pixelization[{instrument}]: regression — vmap log_evidence drifted",
+)
+print(
+    f"  Regression assertion PASSED: log_evidence matches {EXPECTED_LOG_EVIDENCE_HST:.6f}"
+)
diff --git a/likelihood/interferometer/README.md b/likelihood/interferometer/README.md
new file mode 100644
index 0000000..f7d3f4f
--- /dev/null
+++ b/likelihood/interferometer/README.md
@@ -0,0 +1,45 @@
+# likelihood/interferometer
+
+JAX JIT profiling for the PyAutoLens **interferometer** likelihood function — strong-lens reconstruction from radio / sub-millimetre visibility data (e.g. SMA, ALMA, VLA).
+
+The interferometer likelihood path is profiled at **full-pipeline JIT** only, not per-step. The reason is deliberate: the interferometer pipeline runs a Fourier-transformed mapping matrix, a visibilities-space data vector and curvature matrix, and an NNLS solve whose `xp=jnp` threading hasn't been fully characterised yet. Decomposing into per-step JITs would risk missing the cross-step XLA fusion that matters in practice, and would risk hitting library-level JAX blockers that we'd want to raise as separate issues rather than work around in the profiling scripts. Once the full-pipeline JIT is stable across all three model compositions, the per-step breakdown can land as a follow-up.
+
+## What each script measures
+
+1. **Eager baseline** — `FitInterferometer` with `xp=np`, reporting `figure_of_merit` / `log_likelihood`.
+2. **Full-pipeline JIT** — `jax.jit(analysis.log_likelihood_function)` on a pytree-registered `ModelInstance`. Reports lower / compile / first-call / steady-state per-call timings.
+3. **Batched evaluation** — `jax.jit(jax.vmap(full_pipeline))` for the population-of-evaluations regime a sampler actually runs in. Reports per-likelihood cost and speedup vs the single-JIT path.
+4. **Correctness check** — eager ≡ JIT log-likelihood agreement at `rtol=1e-4`.
+5. **Static memory analysis** of the batched program.
+6. **Versioned JSON + PNG write** using the same schema as the imaging scripts so results compare side-by-side.
+
+## Scripts
+
+| Script | Source representation | What it profiles |
+|--------|-----------------------|------------------|
+| [`mge.py`](./mge.py) | MGE (linear light profiles only) | Baseline interferometer path. Exercises the `TuplePrior` pytree support landed in PyAutoFit#1222. |
+| [`pixelization.py`](./pixelization.py) | Rectangular pixelization | Adds source mesh + regularization to the visibility-space inversion. |
+| [`delaunay.py`](./delaunay.py) | Delaunay pixelization | Production-style irregular-source reconstruction in visibility space. |
+
+## Default dataset
+
+`dataset/interferometer/sma/` — an SMA-like mock (pixel scale 0.1″, real-space shape 256×256) committed to this repo. Includes `data.fits`, `noise_map.fits`, `uv_wavelengths.fits`, `positions.json`, and the seeded `tracer.json`. Other instruments (e.g. ALMA, VLA) can be regenerated via the source-of-truth scripts at `autolens_workspace_developer/jax_profiling/dataset_setup/interferometer.py` and copied into `dataset/interferometer/<instrument>/`.
+
+## Headline run-times (populated by Phase 4)
+
+| Script | Instrument | CPU | Laptop GPU | A100 |
+|--------|------------|-----|------------|------|
+| `mge.py` | SMA | _populated_ | _populated_ | _populated_ |
+| `pixelization.py` | SMA | _populated_ | _populated_ | _populated_ |
+| `delaunay.py` | SMA | _populated_ | _populated_ | _populated_ |
+
+Numbers are the **steady-state per-call cost** (single-JIT, post-warmup), in milliseconds. Phase 4's dashboard auto-fills this from the latest `*_summary_v<version>.json` artifacts under `results/likelihood/interferometer/`.
+
+## Output
+
+Each script writes:
+
+```
+results/likelihood/interferometer/<script>_likelihood_summary_<instrument>_v<al.__version__>.json
+results/likelihood/interferometer/<script>_likelihood_summary_<instrument>_v<al.__version__>.png
+```
diff --git a/likelihood/interferometer/delaunay.py b/likelihood/interferometer/delaunay.py
new file mode 100644
index 0000000..32b3e1f
--- /dev/null
+++ b/likelihood/interferometer/delaunay.py
@@ -0,0 +1,1144 @@
+"""
+JAX Profiling: Delaunay Interferometer Likelihood
+=================================================
+
+Profiles the JAX likelihood function for an interferometer dataset where the
+source galaxy is reconstructed using a Delaunay pixelization with cross-
+derivative (``ConstantSplit``) regularization, and the lens galaxy is an
+Isothermal + ExternalShear.
+
+Mirrors ``likelihood/interferometer/pixelization.py`` (Phase 2) with the
+``RectangularUniform`` source replaced by a ``Delaunay`` mesh — matching
+``likelihood/imaging/delaunay.py`` so imaging vs interferometer Delaunay
+results can be compared side-by-side.
+
+Matches the step-by-step pedagogy of ``likelihood/imaging/delaunay.py``
+applied to the visibility-space pipeline. The 11 per-step JIT-profiled stages
+map 1:1 onto sections in
+``autolens_workspace/scripts/interferometer/features/datacube/likelihood_function.py``
+and its single-channel parent
+``interferometer/features/pixelization/likelihood_function.py``.
+
+Pipeline steps (matching the imaging-delaunay numbering for cross-reference;
+the two lens-light steps from the imaging sibling are dropped since the
+interferometer pixelization model has no parametric lens light):
+
+ 1. Ray-trace data grid to source plane.
+ 2. Ray-trace mesh grid (image-plane Overlay vertices) to source plane.
+ 5. Border relocation (data grid + mesh grid).
+ 6. Delaunay triangulation + interpolation + mapper.
+ 7. Mapping matrix.
+ 8. Transformed mapping matrix (NUFFT) — interferometer-specific. Replaces
+    imaging's PSF-convolved blurred mapping matrix; the difference is the
+    Fourier transform to visibility space rather than image-space convolution.
+ 9. Data vector D — visibility-space (real and imaginary components).
+ 10. Curvature matrix F — real and imaginary curvatures summed.
+ 11. Regularization matrix H — ConstantSplit (same as imaging).
+ 12. Reconstruction s = NNLS(F + H, D) (same NNLS path as imaging).
+ 13. Mapped reconstructed visibilities + log evidence (visibility-space χ²).
+
+Measures:
+
+1. Eager baseline: ``FitInterferometer`` with ``xp=np``, print
+   ``figure_of_merit`` / ``log_likelihood``.
+2. Per-step JIT profiling: each pipeline stage above gets its own
+   ``jit_profile()`` call (lower / compile / first-call / steady-state ×10).
+3. Full-pipeline JIT: ``jax.jit(analysis.log_likelihood_function)`` on a
+   pytree-registered ``ModelInstance``. Measure lower / compile / first-call /
+   steady-state per-call.
+4. Batched evaluation (opt-in via ``DELAUNAY_VMAP=1``): ``jax.jit(jax.vmap(...))``.
+   Skipped by default because Delaunay vmap compilation can take 20+ minutes
+   on CPU due to triangulation + interpolation graph size.
+5. Correctness: eager vs JIT log-evidence agreement at ``rtol=1e-4`` for both
+   the per-step recomputation and the full pipeline.
+6. Static memory analysis of the batched program (only when vmap runs).
+7. Results JSON + PNG written to ``results/`` with per-step entries that
+   slot into the same bar-chart shape as ``likelihood/imaging/delaunay.py``.
+
+JIT-blocker notes
+-----------------
+
+Per-step decomposition risks missing cross-step XLA fusion and hitting
+library-level JAX blockers. Caveats from the previous opt-out version that
+still apply:
+
+- ``dataset.transformer.transform_mapping_matrix`` is JIT-friendly for
+  ``TransformerDFT`` (a single matrix multiply) and the default SMA preset
+  uses it. ``TransformerNUFFT`` (pynufft-based) is not JIT-friendly; if you
+  swap the transformer the step-8 timing will fall back to eager-only.
+- The visibility-space χ² in step 13 separates the complex visibilities and
+  noise into real/imag components inside the JIT body (matching the
+  ``pixelization/likelihood_function.py`` reference). Complex-valued JIT
+  with autoarray ``Visibilities`` wrappers is avoided.
+
+Pytree-native parameter inputs
+------------------------------
+
+Uses ``af.ModelInstance`` as the JIT input via PyAutoFit's opt-in pytree
+registration (``autofit.jax.register_model``). Exercises the ``TuplePrior``
+pytree support landed in PyAutoFit#1222.
+"""
+
+import os
+import numpy as np
+import jax
+import jax.numpy as jnp
+import time
+import subprocess
+import sys
+from pathlib import Path
+from contextlib import contextmanager
+
+import autofit as af
+import autolens as al
+from autofit.jax import register_model as _register_model_pytrees
+
+# ---------------------------------------------------------------------------
+# Instrument configuration
+# ---------------------------------------------------------------------------
+
+INSTRUMENTS = {
+    "sma": {"pixel_scale": 0.1, "real_space_shape": (256, 256)},
+    "alma": {"pixel_scale": 0.05, "real_space_shape": (256, 256)},
+}
+
+instrument = "sma"  # <-- change this to profile a different instrument
+
+overlay_shape = (26, 26)
+edge_n_points = 30
+regularization_coefficient = 1.0
+
+
+# ---------------------------------------------------------------------------
+# Profiling helpers
+# ---------------------------------------------------------------------------
+
+class Timer:
+    """Accumulates named timing measurements and prints a summary."""
+
+    def __init__(self):
+        self.records: list[tuple[str, float]] = []
+
+    @contextmanager
+    def section(self, label: str):
+        start = time.perf_counter()
+        yield
+        elapsed = time.perf_counter() - start
+        self.records.append((label, elapsed))
+        print(f"  [{label}] {elapsed:.4f} s")
+
+
+def block(x):
+    """Call block_until_ready if available (JAX arrays)."""
+    if hasattr(x, "block_until_ready"):
+        x.block_until_ready()
+    return x
+
+
+def jit_profile(func, label, *args, n_repeats=10):
+    """JIT-compile *func*, time lower / compile / first call / steady state."""
+    jitted = jax.jit(func)
+
+    with timer.section(f"{label}_lower"):
+        lowered = jitted.lower(*args)
+
+    with timer.section(f"{label}_compile"):
+        compiled = lowered.compile()
+
+    with timer.section(f"{label}_first_call"):
+        result = compiled(*args)
+        block(result)
+
+    with timer.section(f"{label}_steady_x{n_repeats}"):
+        for _ in range(n_repeats):
+            result = compiled(*args)
+            block(result)
+
+    per_call = timer.records[-1][1] / n_repeats
+    print(f"    -> per-call avg: {per_call:.6f} s")
+    return compiled, result
+
+
+timer = Timer()
+likelihood_steps = []  # (label, per_call_seconds) for the final summary
+
+# ===================================================================
+# PART A — Setup (not JIT-compiled)
+# ===================================================================
+
+# ---------------------------------------------------------------------------
+# 1. Dataset
+# ---------------------------------------------------------------------------
+
+print(f"\n--- Dataset loading [{instrument}] ---")
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[1]
+pixel_scale = INSTRUMENTS[instrument]["pixel_scale"]
+real_space_shape = INSTRUMENTS[instrument]["real_space_shape"]
+dataset_path = Path("dataset") / "interferometer" / instrument
+
+if al.util.dataset.should_simulate(str(dataset_path)):
+    raise FileNotFoundError(
+        f"Input dataset missing at '{dataset_path}'. The autolens_profiling "
+        f"repo mirrors only the curated datasets needed for default smoke "
+        f"runs. To regenerate or extend datasets, use the source-of-truth "
+        f"scripts under autolens_workspace_developer/jax_profiling/dataset_setup/, "
+        f"then copy the result into autolens_profiling/dataset/."
+    )
+
+mask_radius = 3.0
+
+real_space_mask = al.Mask2D.circular(
+    shape_native=real_space_shape,
+    pixel_scales=pixel_scale,
+    radius=mask_radius,
+)
+
+with timer.section("dataset_load"):
+    dataset = al.Interferometer.from_fits(
+        data_path=dataset_path / "data.fits",
+        noise_map_path=dataset_path / "noise_map.fits",
+        uv_wavelengths_path=dataset_path / "uv_wavelengths.fits",
+        real_space_mask=real_space_mask,
+        transformer_class=al.TransformerDFT,
+    )
+
+n_visibilities = dataset.uv_wavelengths.shape[0]
+print(f"  Total visibilities: {n_visibilities}")
+
+# ---------------------------------------------------------------------------
+# 2. Image mesh + edge points (Delaunay-specific)
+# ---------------------------------------------------------------------------
+
+print("\n--- Image mesh construction (Delaunay) ---")
+
+with timer.section("image_mesh_overlay"):
+    image_mesh = al.image_mesh.Overlay(shape=overlay_shape)
+    image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=dataset.real_space_mask)
+
+with timer.section("edge_points"):
+    pre_edge_pixels = image_plane_mesh_grid.shape[0]
+    image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(
+        image_plane_mesh_grid=image_plane_mesh_grid,
+        centre=(0.0, 0.0),
+        radius=mask_radius,
+        n_points=edge_n_points,
+    )
+    edge_pixels_total = image_plane_mesh_grid.shape[0] - pre_edge_pixels
+
+n_mesh_vertices = image_plane_mesh_grid.shape[0]
+print(f"  Overlay shape: {overlay_shape}")
+print(f"  Mesh vertices (incl. edge): {n_mesh_vertices}")
+print(f"  Edge points added: {edge_pixels_total}")
+
+# ---------------------------------------------------------------------------
+# 3. Model construction
+# ---------------------------------------------------------------------------
+
+print("\n--- Model construction ---")
+
+with timer.section("model_build"):
+    # GaussianPrior(mean=truth, sigma=small) centres prior-median at the
+    # simulator truth while keeping params free so gradient diagnostics
+    # have dimensionality.
+    mass = af.Model(al.mp.Isothermal)
+    mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.005)
+    mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.005)
+    mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05)
+    _lens_mass_ell = al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0)
+    mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=_lens_mass_ell[0], sigma=0.01)
+    mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=_lens_mass_ell[1], sigma=0.01)
+
+    shear = af.Model(al.mp.ExternalShear)
+    shear.gamma_1 = af.GaussianPrior(mean=0.05, sigma=0.005)
+    shear.gamma_2 = af.GaussianPrior(mean=0.05, sigma=0.005)
+
+    lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)
+
+    mesh = al.mesh.Delaunay(
+        pixels=n_mesh_vertices,
+        zeroed_pixels=edge_pixels_total,
+    )
+    regularization = al.reg.ConstantSplit(coefficient=regularization_coefficient)
+    pixelization = al.Pixelization(mesh=mesh, regularization=regularization)
+
+    source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)
+
+    model = af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+print(f"  Total free parameters: {model.total_free_parameters}")
+print(f"  Delaunay pixels: {n_mesh_vertices}")
+print(f"  Zeroed edge pixels: {edge_pixels_total}")
+
+# ---------------------------------------------------------------------------
+# 4. Instantiate concrete objects from prior medians
+# ---------------------------------------------------------------------------
+
+print("\n--- Instantiate concrete model ---")
+
+with timer.section("instance_from_vector"):
+    param_vector = model.physical_values_from_prior_medians
+    instance = model.instance_from_vector(vector=param_vector)
+
+with timer.section("register_pytrees"):
+    _register_model_pytrees(model)
+
+# JIT input: the instance itself, with all parameter leaves promoted to JAX
+# arrays. The eager NumPy instance is retained for the eager FitInterferometer
+# baseline below.
+params_tree = jax.tree_util.tree_map(jnp.asarray, instance)
+
+tracer = al.Tracer(galaxies=list(instance.galaxies))
+
+# AdaptImages tells FitInterferometer / AnalysisInterferometer where the
+# Delaunay mesh vertices live in the image-plane (separate from the source-
+# plane vertices that get computed by ray-tracing).
+adapt_images = al.AdaptImages(
+    galaxy_image_plane_mesh_grid_dict={
+        instance.galaxies.source: image_plane_mesh_grid,
+    },
+    galaxy_name_image_plane_mesh_grid_dict={
+        "('galaxies', 'source')": image_plane_mesh_grid,
+    },
+)
+
+print(f"  Tracer planes: {tracer.total_planes}")
+
+# ---------------------------------------------------------------------------
+# 5. Configuration summary
+# ---------------------------------------------------------------------------
+
+print("\n--- Configuration (determines run time) ---")
+print(f"  Instrument:              {instrument}")
+print(f"  Pixel scale:             {pixel_scale} arcsec/pixel")
+print(f"  Real-space mask radius:  {mask_radius} arcsec")
+print(f"  Real-space grid shape:   {real_space_shape[0]} x {real_space_shape[1]}")
+print(f"  Visibilities:            {n_visibilities}")
+print(f"  Overlay shape:           {overlay_shape[0]} x {overlay_shape[1]}")
+print(f"  Delaunay vertices:       {n_mesh_vertices}")
+print(f"  Edge zeroed pixels:      {edge_pixels_total}")
+print(f"  Reg. coefficient:        {regularization_coefficient}")
+
+# ---------------------------------------------------------------------------
+# 6. Full-pipeline reference (FitInterferometer) — eager baseline
+# ---------------------------------------------------------------------------
+
+print("\n--- Full FitInterferometer (eager baseline) ---")
+
+with timer.section("fit_interferometer_eager"):
+    fit = al.FitInterferometer(
+        dataset=dataset,
+        tracer=tracer,
+        adapt_images=adapt_images,
+        xp=np,
+    )
+    figure_of_merit_ref = fit.figure_of_merit
+    log_likelihood_ref = fit.log_likelihood
+
+print(f"  figure_of_merit = {figure_of_merit_ref}")
+print(f"  log_likelihood  = {log_likelihood_ref}")
+
+
+# ===================================================================
+# PART B — Per-step JIT profiling
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("PER-STEP JIT PROFILING")
+print("=" * 70)
+
+import autoarray as aa
+
+# Extract raw arrays from autoarray types via .array so they can cross
+# JIT boundaries.  See CLAUDE.md for rationale.
+grid_pix_raw = jnp.array(dataset.grids.pixelization.array)
+mesh_grid_raw = jnp.array(image_plane_mesh_grid.array)
+data_real_jnp = jnp.array(dataset.data.real)
+data_imag_jnp = jnp.array(dataset.data.imag)
+noise_real_jnp = jnp.array(dataset.noise_map.real)
+noise_imag_jnp = jnp.array(dataset.noise_map.imag)
+
+# ---------------------------------------------------------------------------
+# Step 1: Ray-trace data grid to source plane
+# ---------------------------------------------------------------------------
+# Same operation as ``pixelization/likelihood_function.py:__Ray Tracing__``
+# applied to the data grid (one of the two grids the inversion uses).
+
+print("\n--- Step 1: Ray-trace data grid ---")
+
+with timer.section("ray_trace_data_eager"):
+    traced_grids = tracer.traced_grid_2d_list_from(
+        grid=dataset.grids.pixelization, xp=jnp
+    )
+    for tg in traced_grids:
+        block(tg)
+
+print(f"  Number of planes traced: {len(traced_grids)}")
+
+
+def ray_trace_data_raw(grid_raw):
+    grid = aa.Grid2DIrregular(values=grid_raw, xp=jnp)
+    traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp)
+    return jnp.stack([tg.array for tg in traced])
+
+
+_, traced_data_grids_raw = jit_profile(
+    ray_trace_data_raw, "ray_trace_data_jit", grid_pix_raw
+)
+likelihood_steps.append(("Ray-trace data grid", timer.records[-1][1] / 10))
+
+print(f"  traced_data_grids shape: {traced_data_grids_raw.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 2: Ray-trace mesh grid (image-plane vertices) to source plane
+# ---------------------------------------------------------------------------
+# Delaunay-specific: the source-plane mesh vertices are computed in the
+# image-plane via the ``Overlay`` mesh and ray-traced to source-plane. This is
+# the same per-step shape as imaging-delaunay; the underlying tracer call is
+# identical.
+
+print("\n--- Step 2: Ray-trace mesh grid ---")
+
+with timer.section("ray_trace_mesh_eager"):
+    traced_mesh = tracer.traced_grid_2d_list_from(
+        grid=al.Grid2DIrregular(image_plane_mesh_grid), xp=jnp
+    )
+    for tg in traced_mesh:
+        block(tg)
+
+
+def ray_trace_mesh_raw(mesh_raw):
+    grid = aa.Grid2DIrregular(values=mesh_raw, xp=jnp)
+    traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp)
+    return jnp.stack([tg.array for tg in traced])
+
+
+_, traced_mesh_grids_raw = jit_profile(
+    ray_trace_mesh_raw, "ray_trace_mesh_jit", mesh_grid_raw
+)
+likelihood_steps.append(("Ray-trace mesh grid", timer.records[-1][1] / 10))
+
+print(f"  traced_mesh_grids shape: {traced_mesh_grids_raw.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 5: Border relocation (data grid + mesh grid)
+# ---------------------------------------------------------------------------
+# Steps 3-4 from the imaging sibling (lens-light pre-PSF image and PSF-blurred
+# image) don't exist for the interferometer pixelization model — there's no
+# parametric lens light. We jump straight to step 5.
+#
+# Same as ``pixelization/likelihood_function.py:__Border Relocation__``;
+# scipy-based so eager-only.
+
+print("\n--- Step 5: Border relocation ---")
+
+from autoarray.inversion.mesh.border_relocator import BorderRelocator
+
+with timer.section("border_relocator_setup"):
+    border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1)
+
+traced_source_grid = tracer.traced_grid_2d_list_from(
+    grid=dataset.grids.pixelization, xp=jnp
+)[-1]
+traced_mesh_source = tracer.traced_grid_2d_list_from(
+    grid=al.Grid2DIrregular(image_plane_mesh_grid), xp=jnp
+)[-1]
+
+with timer.section("border_relocation_eager"):
+    relocated_grid = border_relocator.relocated_grid_from(grid=traced_source_grid)
+    relocated_mesh_grid = border_relocator.relocated_mesh_grid_from(
+        grid=traced_source_grid,
+        mesh_grid=traced_mesh_source,
+    )
+    block(relocated_grid)
+    block(relocated_mesh_grid)
+
+print(f"  relocated_data_grid shape: {relocated_grid.array.shape}")
+print(f"  relocated_mesh_grid shape: {relocated_mesh_grid.array.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 6: Delaunay triangulation + interpolation + mapper
+# ---------------------------------------------------------------------------
+# scipy-based; same as imaging-delaunay step 6.
+
+print("\n--- Step 6: Delaunay triangulation + Interpolation + Mapper ---")
+
+pixelization_obj = instance.galaxies.source.pixelization
+
+with timer.section("delaunay_interpolation_and_mapper"):
+    interpolator = al.InterpolatorDelaunay(
+        mesh=pixelization_obj.mesh,
+        mesh_grid=relocated_mesh_grid,
+        data_grid=relocated_grid,
+    )
+    mapper = al.Mapper(
+        interpolator=interpolator,
+        image_plane_mesh_grid=image_plane_mesh_grid,
+        xp=jnp,
+    )
+
+print(f"  mapper.pixels (source): {mapper.pixels}")
+print(f"  pix_indexes shape: {mapper.pix_indexes_for_sub_slim_index.shape}")
+
+# ---------------------------------------------------------------------------
+# Steps 7-13: Extract matrices from FitInterferometer.inversion for consistency
+# ---------------------------------------------------------------------------
+# The FitInterferometer pipeline handles edge-pixel zeroing, curvature-diagonal
+# adjustments, and settings that are difficult to replicate manually. We extract
+# the production matrices upfront and then JIT-profile the linear-algebra ops
+# on the reference inputs, matching the imaging-sibling pattern.
+
+print("\n--- Extracting inversion matrices from FitInterferometer ---")
+
+inversion = fit.inversion
+
+with timer.section("extract_inversion_matrices"):
+    # ``operated_mapping_matrix`` is the NUFFT-transformed mapping matrix
+    # (complex-valued, shape (n_vis, source_pixels)). Imaging's equivalent is
+    # the PSF-convolved blurred mapping matrix.
+    transformed_mm_ref = jnp.asarray(inversion.operated_mapping_matrix)
+    mapping_matrix_ref = jnp.asarray(inversion.mapping_matrix)
+
+    inv_mapper = inversion.cls_list_from(cls=al.Mapper)[0]
+    neighbors = inv_mapper.neighbors
+    neighbors_array = jnp.array(np.asarray(neighbors))
+    neighbors_sizes = jnp.array(neighbors.sizes)
+
+print(f"  transformed_mapping_matrix shape: {transformed_mm_ref.shape}")
+print(f"  transformed_mapping_matrix dtype: {transformed_mm_ref.dtype}")
+print(f"  mapping_matrix shape: {mapping_matrix_ref.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 7: Mapping matrix
+# ---------------------------------------------------------------------------
+# Same as ``pixelization/likelihood_function.py:__Mapping Matrix__``.
+
+print("\n--- Step 7: Mapping matrix ---")
+
+with timer.section("mapping_matrix"):
+    mapping_matrix = inv_mapper.mapping_matrix
+
+print(f"  mapping_matrix shape: {mapping_matrix.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 8: Transformed mapping matrix (NUFFT) — interferometer-specific
+# ---------------------------------------------------------------------------
+# Replaces the imaging sibling's "Blurred mapping matrix (PSF convolution)".
+#
+# For ``TransformerDFT`` this is a single complex matrix multiply
+# ``D @ M`` where D is the discrete Fourier matrix (n_vis × n_image) and
+# M is the real mapping matrix (n_image × source_pixels). JIT-friendly.
+# For ``TransformerNUFFT`` (pynufft-based, ALMA-scale) this step would fall
+# back to eager — flag if you swap the transformer.
+
+print("\n--- Step 8: Transformed mapping matrix (NUFFT) ---")
+
+with timer.section("transformed_mapping_matrix_eager"):
+    transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(
+        mapping_matrix=mapping_matrix_ref
+    )
+    block(transformed_mapping_matrix)
+
+
+def compute_transformed_mapping_matrix(mapping_matrix):
+    return dataset.transformer.transform_mapping_matrix(mapping_matrix=mapping_matrix)
+
+
+# JIT-profile the full inversion setup pipeline (steps 5-8 combined) from a
+# pytree ModelInstance. This is the cube-relevant per-channel cost: per
+# AnalysisFactor the NUFFT mapping-matrix construction has to rerun because
+# uv_wavelengths is channel-specific.
+def transformed_mm_from_params(params_tree):
+    """Inversion setup from a pytree ModelInstance — full chain through NUFFT."""
+    t = al.Tracer(galaxies=list(params_tree.galaxies))
+    adapt_images_jax = al.AdaptImages(
+        galaxy_image_plane_mesh_grid_dict={
+            params_tree.galaxies.source: image_plane_mesh_grid,
+        },
+        galaxy_name_image_plane_mesh_grid_dict={
+            "('galaxies', 'source')": image_plane_mesh_grid,
+        },
+    )
+    fit_jax = al.FitInterferometer(
+        dataset=dataset,
+        tracer=t,
+        adapt_images=adapt_images_jax,
+        xp=jnp,
+    )
+    return jnp.asarray(fit_jax.inversion.operated_mapping_matrix)
+
+
+_, transformed_mm_jit = jit_profile(
+    transformed_mm_from_params, "inversion_setup_jit", params_tree
+)
+likelihood_steps.append(
+    ("Inversion setup (steps 5-8 combined, incl. NUFFT)", timer.records[-1][1] / 10)
+)
+
+print(f"  transformed_mapping_matrix (JIT) shape: {transformed_mm_jit.shape}")
+print(f"  transformed_mapping_matrix (JIT) dtype: {transformed_mm_jit.dtype}")
+
+# Use the reference matrices for the linear-algebra steps below.
+transformed_mm_real_jnp = jnp.real(transformed_mm_ref)
+transformed_mm_imag_jnp = jnp.imag(transformed_mm_ref)
+
+# ---------------------------------------------------------------------------
+# Step 9: Data vector (D) — visibility-space
+# ---------------------------------------------------------------------------
+# Same as ``pixelization/likelihood_function.py:__Data Vector (D)__``.
+
+print("\n--- Step 9: Data vector (D) ---")
+
+
+def compute_data_vector(
+    transformed_mm_real, transformed_mm_imag, data_real, data_imag,
+    noise_real, noise_imag,
+):
+    # Visibility-space data vector: D_i = sum_j f_ij d_j / sigma_j^2 (real + imag).
+    weighted_data_real = data_real / (noise_real ** 2)
+    weighted_data_imag = data_imag / (noise_imag ** 2)
+    return jnp.matmul(transformed_mm_real.T, weighted_data_real) + jnp.matmul(
+        transformed_mm_imag.T, weighted_data_imag
+    )
+
+
+with timer.section("data_vector_eager"):
+    data_vector = compute_data_vector(
+        transformed_mm_real_jnp, transformed_mm_imag_jnp,
+        data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp,
+    )
+    block(data_vector)
+
+_, data_vector = jit_profile(
+    compute_data_vector, "data_vector_jit",
+    transformed_mm_real_jnp, transformed_mm_imag_jnp,
+    data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp,
+)
+likelihood_steps.append(("Data vector (D)", timer.records[-1][1] / 10))
+
+print(f"  data_vector shape: {data_vector.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 10: Curvature matrix (F)
+# ---------------------------------------------------------------------------
+# Same as ``pixelization/likelihood_function.py:__Curvature Matrix (F)__``:
+# F = sum_j f_ij f_kj / sigma_j^2, computed separately for real / imag and summed.
+
+print("\n--- Step 10: Curvature matrix (F) ---")
+
+no_reg_list = list(inversion.no_regularization_index_list)
+
+
+def compute_curvature_matrix(
+    transformed_mm_real, transformed_mm_imag, noise_real, noise_imag,
+):
+    real_curv = al.util.inversion.curvature_matrix_via_mapping_matrix_from(
+        mapping_matrix=transformed_mm_real,
+        noise_map=noise_real,
+        settings=fit.settings,
+        add_to_curvature_diag=True,
+        no_regularization_index_list=no_reg_list,
+        xp=jnp,
+    )
+    imag_curv = al.util.inversion.curvature_matrix_via_mapping_matrix_from(
+        mapping_matrix=transformed_mm_imag,
+        noise_map=noise_imag,
+        settings=fit.settings,
+        add_to_curvature_diag=False,
+        no_regularization_index_list=no_reg_list,
+        xp=jnp,
+    )
+    return real_curv + imag_curv
+
+
+with timer.section("curvature_matrix_eager"):
+    curvature_matrix = compute_curvature_matrix(
+        transformed_mm_real_jnp, transformed_mm_imag_jnp, noise_real_jnp, noise_imag_jnp,
+    )
+    block(curvature_matrix)
+
+_, curvature_matrix = jit_profile(
+    compute_curvature_matrix, "curvature_matrix_jit",
+    transformed_mm_real_jnp, transformed_mm_imag_jnp, noise_real_jnp, noise_imag_jnp,
+)
+likelihood_steps.append(("Curvature matrix (F)", timer.records[-1][1] / 10))
+
+print(f"  curvature_matrix shape: {curvature_matrix.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 11: Regularization matrix (H) — ConstantSplit scheme
+# ---------------------------------------------------------------------------
+# Same as ``pixelization/likelihood_function.py:__Regularization Matrix (H)__``;
+# extracted from the inversion for consistency with the production setup.
+
+print("\n--- Step 11: Regularization matrix (ConstantSplit) ---")
+
+with timer.section("regularization_matrix_eager"):
+    regularization_matrix = jnp.array(inversion.regularization_matrix)
+    block(regularization_matrix)
+
+likelihood_steps.append(("Regularization matrix (H)", timer.records[-1][1]))
+
+print(f"  regularization_matrix shape: {regularization_matrix.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 12: Regularized reconstruction: s = NNLS(F + H, D)
+# ---------------------------------------------------------------------------
+# Same NNLS path as the imaging sibling. For well-conditioned ConstantSplit at
+# the prior median this reduces to a linear solve (no negative source pixels),
+# but we use the NNLS solver to match the production AnalysisInterferometer
+# behaviour exactly.
+
+print("\n--- Step 12: Regularized reconstruction ---")
+
+
+def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix):
+    curvature_reg_matrix = curvature_matrix + regularization_matrix
+    return al.util.inversion.reconstruction_positive_only_from(
+        data_vector=data_vector,
+        curvature_reg_matrix=curvature_reg_matrix,
+        xp=jnp,
+    )
+
+
+with timer.section("reconstruction_eager"):
+    reconstruction = compute_reconstruction(
+        jnp.array(data_vector),
+        jnp.array(curvature_matrix),
+        jnp.array(regularization_matrix),
+    )
+    block(reconstruction)
+
+_, reconstruction = jit_profile(
+    compute_reconstruction, "reconstruction_jit",
+    jnp.array(data_vector), jnp.array(curvature_matrix), jnp.array(regularization_matrix),
+)
+likelihood_steps.append(("Regularized reconstruction", timer.records[-1][1] / 10))
+
+print(f"  reconstruction shape: {reconstruction.shape}")
+
+# ---------------------------------------------------------------------------
+# Step 13: Mapped reconstructed visibilities + log evidence
+# ---------------------------------------------------------------------------
+# Same five-term log-evidence as the imaging sibling but with visibility-space
+# χ² (real + imag). Matches the formula in
+# ``pixelization/likelihood_function.py:__Likelihood Function — Five Terms__``.
+
+print("\n--- Step 13: Mapped reconstructed visibilities + log evidence ---")
+
+
+def compute_log_evidence(
+    data_real, data_imag, noise_real, noise_imag,
+    transformed_mm_real, transformed_mm_imag,
+    reconstruction, curvature_matrix, regularization_matrix, mapper_indices,
+):
+    """Visibility-space log-evidence — five-term formula matching the production
+    ``FitInterferometer.log_evidence``.
+
+    -2 ln e = chi^2 + s^T H s + ln[det(F+H)] - ln[det(H)] + noise_norm
+    """
+    # Mapped reconstructed visibilities (real / imag separately)
+    mapped_real = jnp.matmul(transformed_mm_real, reconstruction)
+    mapped_imag = jnp.matmul(transformed_mm_imag, reconstruction)
+
+    # χ² in visibility space (real + imag)
+    chi_real = jnp.sum(((data_real - mapped_real) / noise_real) ** 2)
+    chi_imag = jnp.sum(((data_imag - mapped_imag) / noise_imag) ** 2)
+    chi_squared = chi_real + chi_imag
+
+    # s^T H s
+    regularization_term = jnp.dot(
+        reconstruction, jnp.dot(regularization_matrix, reconstruction)
+    )
+
+    # Complexity terms (Cholesky log-det matching production)
+    curvature_reg_matrix = curvature_matrix + regularization_matrix
+    creg_reduced = curvature_reg_matrix[mapper_indices][:, mapper_indices]
+    reg_reduced = regularization_matrix[mapper_indices][:, mapper_indices]
+    log_det_curvature_reg = 2.0 * jnp.sum(
+        jnp.log(jnp.diag(jnp.linalg.cholesky(creg_reduced)))
+    )
+    log_det_regularization = 2.0 * jnp.sum(
+        jnp.log(jnp.diag(jnp.linalg.cholesky(reg_reduced)))
+    )
+
+    # Noise normalisation (real + imag)
+    noise_normalization = (
+        jnp.sum(jnp.log(2 * jnp.pi * noise_real ** 2))
+        + jnp.sum(jnp.log(2 * jnp.pi * noise_imag ** 2))
+    )
+
+    return -0.5 * (
+        chi_squared + regularization_term + log_det_curvature_reg
+        - log_det_regularization + noise_normalization
+    )
+
+
+# For the JIT correctness check we recompute log_evidence using the inversion's
+# own reconstruction and curvature matrix to avoid accumulated FP drift
+# (matching the imaging sibling's pattern).
+mapper_indices_jnp = jnp.array(np.asarray(inversion.mapper_indices))
+inv_recon_jnp = jnp.asarray(inversion.reconstruction)
+inv_curv_jnp = jnp.asarray(inversion.curvature_matrix)
+reg_jnp = jnp.array(regularization_matrix)
+
+with timer.section("log_evidence_eager"):
+    log_evidence = compute_log_evidence(
+        data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp,
+        transformed_mm_real_jnp, transformed_mm_imag_jnp,
+        reconstruction, curvature_matrix, reg_jnp, mapper_indices_jnp,
+    )
+    block(log_evidence)
+
+_, log_evidence = jit_profile(
+    compute_log_evidence, "log_evidence_jit",
+    data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp,
+    transformed_mm_real_jnp, transformed_mm_imag_jnp,
+    reconstruction, curvature_matrix, reg_jnp, mapper_indices_jnp,
+)
+likelihood_steps.append(("Mapped recon + log evidence", timer.records[-1][1] / 10))
+
+print(f"  log_evidence (step-by-step) = {log_evidence}")
+
+# Correctness check: use the inversion's own reconstruction and curvature matrix
+log_evidence_check = compute_log_evidence(
+    data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp,
+    transformed_mm_real_jnp, transformed_mm_imag_jnp,
+    inv_recon_jnp, inv_curv_jnp, reg_jnp, mapper_indices_jnp,
+)
+print(f"  log_evidence (inv matrices) = {log_evidence_check}")
+print(f"  log_evidence (reference)    = {figure_of_merit_ref}")
+
+np.testing.assert_allclose(
+    float(log_evidence_check),
+    float(figure_of_merit_ref),
+    rtol=1e-4,
+    err_msg=(
+        "Per-step log_evidence from inversion matrices does not match "
+        "FitInterferometer.log_evidence"
+    ),
+)
+print(
+    "  Assertion PASSED: inversion-matrix log_evidence matches "
+    "FitInterferometer.log_evidence"
+)
+
+
+# ===================================================================
+# PART C — Full-pipeline JIT (for comparison)
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("FULL-PIPELINE JIT")
+print("=" * 70)
+
+analysis = al.AnalysisInterferometer(
+    dataset=dataset, adapt_images=adapt_images, use_jax=True
+)
+
+def full_pipeline_from_params(params_tree):
+    """Full interferometer likelihood from a pytree-shaped ``ModelInstance``.
+
+    No flat-vector unpacking inside the trace — the instance crosses the JIT
+    boundary directly, with constants (redshifts, etc.) kept static via the
+    ``aux_data`` partition set up by ``autofit.jax.register_model``.
+    """
+    return analysis.log_likelihood_function(instance=params_tree)
+
+_, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree)
+full_pipeline_per_call = timer.records[-1][1] / 10
+
+print(f"  full log_evidence = {full_result}")
+
+# Correctness: for inversion models (pixelization + regularization), the
+# analysis "log_likelihood_function" actually returns the log-evidence
+# (= figure_of_merit), which includes the regularization/determinant terms.
+# Match against figure_of_merit_ref, not log_likelihood_ref.
+np.testing.assert_allclose(
+    float(full_result),
+    float(figure_of_merit_ref),
+    rtol=1e-4,
+    err_msg="interferometer/delaunay: JIT log-evidence does not match eager figure_of_merit",
+)
+print("  Eager-vs-JIT correctness PASSED")
+
+# ===================================================================
+# PART D — vmap (opt-in) + correctness
+# ===================================================================
+#
+# Delaunay vmap compilation can take 20+ minutes on CPU due to the size of
+# the triangulation + interpolation XLA graph. Skipped by default — set
+# DELAUNAY_VMAP=1 to opt in.
+
+print("\n--- vmap batched evaluation ---")
+
+run_vmap = os.environ.get("DELAUNAY_VMAP", "0") == "1"
+
+batch_size = 3
+vmap_batch_time = None
+vmap_per_call = None
+vmap_speedup = None
+result_vmap = None
+vmapped_full = None
+parameters = None
+
+_n_leaves = len(jax.tree_util.tree_leaves(params_tree))
+if not run_vmap:
+    print("  SKIPPED: opt-in via DELAUNAY_VMAP=1 (compilation can take 20+ minutes).")
+elif _n_leaves == 0:
+    print(f"  SKIPPED: model has 0 free parameters (all fixed to truth); "
+          f"vmap requires at least one array leaf.")
+else:
+    parameters = jax.tree_util.tree_map(
+        lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)),
+        params_tree,
+    )
+
+    vmapped_full = jax.jit(jax.vmap(full_pipeline_from_params))
+
+    with timer.section("vmap_first_call"):
+        result_vmap = vmapped_full(parameters)
+        block(result_vmap)
+
+    n_vmap_repeats = 10
+    with timer.section(f"vmap_steady_x{n_vmap_repeats}"):
+        for _ in range(n_vmap_repeats):
+            result_vmap = vmapped_full(parameters)
+            block(result_vmap)
+
+    vmap_batch_time = timer.records[-1][1] / n_vmap_repeats
+    vmap_per_call = vmap_batch_time / batch_size
+    vmap_speedup = full_pipeline_per_call / vmap_per_call
+
+    print(f"  batch results = {result_vmap}")
+    print(f"  vmap batch of {batch_size}:   {vmap_batch_time:.6f} s")
+    print(f"  vmap per call:         {vmap_per_call:.6f} s")
+    print(f"  single JIT per call:   {full_pipeline_per_call:.6f} s")
+    print(f"  vmap speedup:          {vmap_speedup:.1f}x faster per likelihood")
+
+    np.testing.assert_allclose(
+        np.array(result_vmap),
+        float(full_result),
+        rtol=1e-4,
+        err_msg="interferometer/delaunay: JAX vmap likelihood mismatch",
+    )
+    print("  vmap-vs-single-JIT correctness PASSED")
+
+# ===================================================================
+# PART E — Static memory analysis (only if vmap ran)
+# ===================================================================
+
+print("\n--- Static memory analysis ---")
+
+if vmapped_full is None:
+    print("  SKIPPED: vmap path was not exercised this run.")
+    memory_analysis = None
+else:
+    lowered_batched = vmapped_full.lower(parameters)
+    compiled_batched = lowered_batched.compile()
+
+    memory_analysis = compiled_batched.memory_analysis()
+    print(f"  Output size:  {memory_analysis.output_size_in_bytes / 1024**2:.3f} MB")
+    print(f"  Temp size:    {memory_analysis.temp_size_in_bytes / 1024**2:.3f} MB")
+    print(
+        f"  Total:        "
+        f"{(memory_analysis.output_size_in_bytes + memory_analysis.temp_size_in_bytes) / 1024**2:.3f} MB"
+    )
+
+# ===================================================================
+# JAX Likelihood Function Summary + artefacts
+# ===================================================================
+
+import json
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+al_version = al.__version__
+
+print("\n" + "=" * 70)
+print(f"JAX LIKELIHOOD FUNCTION SUMMARY — {instrument.upper()} — v{al_version}")
+print("=" * 70)
+print(f"  Instrument:              {instrument}")
+print(f"  Pixel scale:             {pixel_scale} arcsec/pixel")
+print(f"  Real-space mask radius:  {mask_radius} arcsec")
+print(f"  Real-space grid shape:   {real_space_shape[0]} x {real_space_shape[1]}")
+print(f"  Visibilities:            {n_visibilities}")
+print(f"  Delaunay vertices:       {n_mesh_vertices}")
+print(f"  Edge zeroed pixels:      {edge_pixels_total}")
+print("-" * 70)
+print(f"  Eager log_likelihood:    {log_likelihood_ref}")
+print(f"  Eager figure_of_merit:   {figure_of_merit_ref}  (log-evidence)")
+print(f"  JIT  log-evidence:       {float(full_result)}")
+print("-" * 70)
+
+max_label = max(len(label) for label, _ in likelihood_steps)
+step_total = 0.0
+for i, (label, per_call) in enumerate(likelihood_steps, 1):
+    print(f"  {i:>2}. {label:<{max_label}}  {per_call:>12.6f} s")
+    step_total += per_call
+
+print("-" * 70)
+print(f"      {'TOTAL (step-by-step)':<{max_label}}  {step_total:>12.6f} s")
+print(f"      {'Full pipeline (single JIT)':<{max_label}}  {full_pipeline_per_call:>12.6f} s")
+if vmap_per_call is not None:
+    print(f"      {'vmap batch (per call)':<{max_label}}  {vmap_per_call:>12.6f} s")
+    print(f"      {'vmap speedup vs single JIT':<{max_label}}  {vmap_speedup:>11.1f}x")
+else:
+    print(f"      {'vmap':<{max_label}}  {'SKIPPED':>12}")
+print("=" * 70)
+
+# --- Save results dictionary ---
+
+if vmap_per_call is None:
+    if not run_vmap:
+        vmap_payload = "SKIPPED — opt-in via DELAUNAY_VMAP=1"
+    else:
+        vmap_payload = "SKIPPED — model has 0 free parameters (all fixed to truth)"
+else:
+    vmap_payload = {
+        "batch_size": batch_size,
+        "batch_time": vmap_batch_time,
+        "per_call": vmap_per_call,
+        "speedup_vs_single_jit": round(vmap_speedup, 1),
+    }
+
+likelihood_summary = {
+    "autolens_version": al_version,
+    "instrument": instrument,
+    "model": "delaunay",
+    "configuration": {
+        "pixel_scale_arcsec": pixel_scale,
+        "mask_radius_arcsec": mask_radius,
+        "real_space_shape": list(real_space_shape),
+        "visibilities": int(n_visibilities),
+        "overlay_shape": list(overlay_shape),
+        "edge_n_points": edge_n_points,
+        "delaunay_vertices": int(n_mesh_vertices),
+        "edge_zeroed_pixels": int(edge_pixels_total),
+        "regularization_coefficient": regularization_coefficient,
+    },
+    "log_likelihood_eager": float(log_likelihood_ref),
+    "figure_of_merit_eager": float(figure_of_merit_ref),
+    "log_evidence_jit": float(full_result),
+    "steps": {label: per_call for label, per_call in likelihood_steps},
+    "total_step_by_step": step_total,
+    "full_pipeline_single_jit": full_pipeline_per_call,
+    "vmap": vmap_payload,
+    "memory_mb": None if memory_analysis is None else {
+        "output": memory_analysis.output_size_in_bytes / 1024**2,
+        "temp": memory_analysis.temp_size_in_bytes / 1024**2,
+    },
+}
+
+results_dir = _workspace_root / "results" / "likelihood" / "interferometer"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+dict_path = results_dir / f"delaunay_likelihood_summary_{instrument}_v{al_version}.json"
+dict_path.write_text(json.dumps(likelihood_summary, indent=2))
+print(f"\n  Results dict saved to: {dict_path}")
+
+# --- Save bar chart ---
+
+labels = [label for label, _ in likelihood_steps]
+times = [per_call for _, per_call in likelihood_steps]
+
+fig, ax = plt.subplots(figsize=(10, 6))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color="#4C72B0", edgecolor="white", height=0.6)
+
+for bar, t in zip(bars, times):
+    ax.text(
+        bar.get_width() + max(times) * 0.01,
+        bar.get_y() + bar.get_height() / 2,
+        f"{t:.6f} s",
+        va="center",
+        fontsize=9,
+    )
+
+ax.axvline(
+    full_pipeline_per_call,
+    color="#C44E52",
+    linestyle="--",
+    linewidth=1.5,
+    label=f"Full pipeline (single JIT): {full_pipeline_per_call:.6f} s",
+)
+if vmap_per_call is not None:
+    ax.axvline(
+        vmap_per_call,
+        color="#55A868",
+        linestyle="--",
+        linewidth=1.5,
+        label=f"vmap batch per call: {vmap_per_call:.6f} s ({vmap_speedup:.1f}x faster)",
+    )
+
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=10)
+ax.invert_yaxis()
+ax.set_xlabel("Time per call (s)", fontsize=11)
+fig.suptitle(
+    f"Delaunay Interferometer Likelihood — {instrument.upper()}",
+    fontsize=12,
+    fontweight="bold",
+)
+ax.set_title(
+    f"AutoLens v{al_version}  |  {pixel_scale}\"/px  |  "
+    f"{real_space_shape[0]}x{real_space_shape[1]} real-space  |  "
+    f"{n_visibilities} visibilities  |  {n_mesh_vertices} Delaunay verts  |  "
+    f"total: {step_total:.6f} s",
+    fontsize=9,
+)
+ax.legend(loc="lower right", fontsize=9)
+ax.margins(x=0.15)
+fig.tight_layout()
+
+chart_path = results_dir / f"delaunay_likelihood_summary_{instrument}_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to:    {chart_path}")
+
+
+# ===================================================================
+# Regression assertion — realistic-scale deterministic log-evidence
+# ===================================================================
+#
+# Simulator truth parameters via GaussianPrior(mean=truth, sigma=small)
+# make the full-pipeline log-evidence deterministic at the prior median.
+EXPECTED_LOG_EVIDENCE_SMA = -3167.5258928840763
+
+if EXPECTED_LOG_EVIDENCE_SMA is None:
+    print(
+        f"\n  Regression assertion SKIPPED — "
+        f"capture this run's eager log_evidence ({figure_of_merit_ref}) "
+        f"and set EXPECTED_LOG_EVIDENCE_SMA in this script."
+    )
+else:
+    np.testing.assert_allclose(
+        figure_of_merit_ref,
+        EXPECTED_LOG_EVIDENCE_SMA,
+        rtol=1e-4,
+        err_msg=(
+            f"interferometer/delaunay[{instrument}]: regression — eager log_evidence "
+            f"drifted (got {figure_of_merit_ref}, expected {EXPECTED_LOG_EVIDENCE_SMA})"
+        ),
+    )
+    print(
+        f"  Eager regression assertion PASSED: log_evidence matches "
+        f"{EXPECTED_LOG_EVIDENCE_SMA:.6f}"
+    )
+    np.testing.assert_allclose(
+        float(full_result),
+        EXPECTED_LOG_EVIDENCE_SMA,
+        rtol=1e-4,
+        err_msg=f"interferometer/delaunay[{instrument}]: regression — full log_evidence drifted",
+    )
+    print(f"  Full-pipeline regression assertion PASSED")
+    if result_vmap is not None:
+        np.testing.assert_allclose(
+            np.array(result_vmap),
+            EXPECTED_LOG_EVIDENCE_SMA,
+            rtol=1e-4,
+            err_msg=f"interferometer/delaunay[{instrument}]: regression — vmap log_evidence drifted",
+        )
+        print(f"  vmap regression assertion PASSED")
diff --git a/likelihood/interferometer/mge.py b/likelihood/interferometer/mge.py
new file mode 100644
index 0000000..0710e38
--- /dev/null
+++ b/likelihood/interferometer/mge.py
@@ -0,0 +1,509 @@
+"""
+JAX Profiling: MGE Interferometer Likelihood
+=============================================
+
+Profiles the JAX likelihood function for an interferometer dataset where the
+source galaxy's light is modelled with a multi-Gaussian expansion (MGE) and
+the lens galaxy is an Isothermal + ExternalShear.
+
+Unlike the imaging MGE profiling script, this script deliberately does not
+per-step JIT the inversion pipeline. The interferometer path exercises a
+Fourier-transformed mapping matrix, a visibilities-space data vector /
+curvature matrix, and an NNLS solve whose ``xp=jnp`` threading has not been
+fully characterised. Per-step decomposition risks missing cross-step XLA
+fusion that matters in practice, and risks hitting library-level JAX
+blockers that we would want to raise as separate issues rather than work
+around here. Once the full-pipeline JIT is stable on interferometer, the
+per-step breakdown can land as a follow-up.
+
+Instead, this script measures:
+
+1. Eager baseline: ``FitInterferometer`` with ``xp=np``, print
+   ``figure_of_merit`` / ``log_likelihood``.
+2. Full-pipeline JIT: ``jax.jit(analysis.log_likelihood_function)`` on a
+   pytree-registered ``ModelInstance``. Measure lower / compile / first-call /
+   steady-state per-call.
+3. Batched evaluation: ``jax.jit(jax.vmap(full_pipeline))``. Measure
+   per-likelihood cost and speedup vs the single-JIT path.
+4. Correctness: eager vs JIT log-likelihood agreement at ``rtol=1e-4``.
+5. Static memory analysis of the batched program.
+6. Results JSON + PNG written to ``results/`` using the same schema as the
+   imaging profiling scripts so they can be compared side-by-side.
+
+Pytree-native parameter inputs
+------------------------------
+
+Uses ``af.ModelInstance`` as the JIT input via PyAutoFit's opt-in pytree
+registration (``autofit.jax.register_model``). This matches the pattern in
+``likelihood/imaging/mge.py`` and exercises the ``TuplePrior`` pytree
+support landed in PyAutoFit#1222.
+"""
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import time
+import subprocess
+import sys
+from pathlib import Path
+from contextlib import contextmanager
+
+import autofit as af
+import autolens as al
+from autofit.jax import register_model as _register_model_pytrees
+
+# ---------------------------------------------------------------------------
+# Instrument configuration
+# ---------------------------------------------------------------------------
+
+INSTRUMENTS = {
+    "sma": {"pixel_scale": 0.1, "real_space_shape": (256, 256)},
+    "alma": {"pixel_scale": 0.05, "real_space_shape": (256, 256)},
+}
+
+instrument = "sma"  # <-- change this to profile a different instrument
+
+
+# ---------------------------------------------------------------------------
+# Profiling helpers (copied verbatim from imaging/mge.py)
+# ---------------------------------------------------------------------------
+
+
+class Timer:
+    """Accumulates named timing measurements and prints a summary."""
+
+    def __init__(self):
+        self.records: list[tuple[str, float]] = []
+
+    @contextmanager
+    def section(self, label: str):
+        start = time.perf_counter()
+        yield
+        elapsed = time.perf_counter() - start
+        self.records.append((label, elapsed))
+        print(f"  [{label}] {elapsed:.4f} s")
+
+
+def block(x):
+    """Call block_until_ready if available (JAX arrays)."""
+    if hasattr(x, "block_until_ready"):
+        x.block_until_ready()
+    return x
+
+
+def jit_profile(func, label, *args, n_repeats=10):
+    """JIT-compile *func*, time lower / compile / first call / steady state."""
+    jitted = jax.jit(func)
+
+    with timer.section(f"{label}_lower"):
+        lowered = jitted.lower(*args)
+
+    with timer.section(f"{label}_compile"):
+        compiled = lowered.compile()
+
+    with timer.section(f"{label}_first_call"):
+        result = compiled(*args)
+        block(result)
+
+    with timer.section(f"{label}_steady_x{n_repeats}"):
+        for _ in range(n_repeats):
+            result = compiled(*args)
+            block(result)
+
+    per_call = timer.records[-1][1] / n_repeats
+    print(f"    -> per-call avg: {per_call:.6f} s")
+    return compiled, result
+
+
+timer = Timer()
+
+# ===================================================================
+# PART A — Setup (not JIT-compiled)
+# ===================================================================
+
+# ---------------------------------------------------------------------------
+# 1. Dataset
+# ---------------------------------------------------------------------------
+
+print(f"\n--- Dataset loading [{instrument}] ---")
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[1]
+pixel_scale = INSTRUMENTS[instrument]["pixel_scale"]
+real_space_shape = INSTRUMENTS[instrument]["real_space_shape"]
+dataset_path = Path("dataset") / "interferometer" / instrument
+
+if al.util.dataset.should_simulate(str(dataset_path)):
+    raise FileNotFoundError(
+        f"Input dataset missing at '{dataset_path}'. The autolens_profiling "
+        f"repo mirrors only the curated datasets needed for default smoke "
+        f"runs. To regenerate or extend datasets, use the source-of-truth "
+        f"scripts under autolens_workspace_developer/jax_profiling/dataset_setup/, "
+        f"then copy the result into autolens_profiling/dataset/."
+    )
+
+mask_radius = 3.0
+
+real_space_mask = al.Mask2D.circular(
+    shape_native=real_space_shape,
+    pixel_scales=pixel_scale,
+    radius=mask_radius,
+)
+
+with timer.section("dataset_load"):
+    dataset = al.Interferometer.from_fits(
+        data_path=dataset_path / "data.fits",
+        noise_map_path=dataset_path / "noise_map.fits",
+        uv_wavelengths_path=dataset_path / "uv_wavelengths.fits",
+        real_space_mask=real_space_mask,
+        transformer_class=al.TransformerDFT,
+    )
+
+n_visibilities = dataset.uv_wavelengths.shape[0]
+print(f"  Total visibilities: {n_visibilities}")
+
+# ---------------------------------------------------------------------------
+# 2. Model construction
+# ---------------------------------------------------------------------------
+
+print("\n--- Model construction ---")
+
+with timer.section("model_build"):
+    # GaussianPrior(mean=truth, sigma=small) centres prior-median at the
+    # simulator truth while keeping params free so gradient diagnostics
+    # have dimensionality.
+    mass = af.Model(al.mp.Isothermal)
+    mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.005)
+    mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.005)
+    mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05)
+    _lens_mass_ell = al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0)
+    mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=_lens_mass_ell[0], sigma=0.01)
+    mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=_lens_mass_ell[1], sigma=0.01)
+
+    shear = af.Model(al.mp.ExternalShear)
+    shear.gamma_1 = af.GaussianPrior(mean=0.05, sigma=0.005)
+    shear.gamma_2 = af.GaussianPrior(mean=0.05, sigma=0.005)
+
+    lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)
+
+    # Simulator truth source centre is (0.1, 0.1); set via mge_model_from's
+    # centre kwarg so the shared centre prior's median lands there.
+    source_bulge = al.model_util.mge_model_from(
+        mask_radius=mask_radius,
+        total_gaussians=20,
+        centre_prior_is_uniform=False,
+        centre=(0.1, 0.1),
+        centre_sigma=0.005,
+    )
+
+    source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)
+
+    model = af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+print(f"  Total free parameters: {model.total_free_parameters}")
+
+# ---------------------------------------------------------------------------
+# 3. Instantiate concrete objects from prior medians
+# ---------------------------------------------------------------------------
+
+print("\n--- Instantiate concrete model ---")
+
+with timer.section("instance_from_vector"):
+    param_vector = model.physical_values_from_prior_medians
+    instance = model.instance_from_vector(vector=param_vector)
+
+with timer.section("register_pytrees"):
+    _register_model_pytrees(model)
+
+# JIT input: the instance itself, with all parameter leaves promoted to JAX
+# arrays. The eager NumPy instance is retained for the eager FitInterferometer
+# baseline below.
+params_tree = jax.tree_util.tree_map(jnp.asarray, instance)
+
+tracer = al.Tracer(galaxies=list(instance.galaxies))
+
+print(f"  Tracer planes: {tracer.total_planes}")
+
+# ---------------------------------------------------------------------------
+# 4. Configuration summary
+# ---------------------------------------------------------------------------
+
+from autogalaxy.profiles.basis import Basis as _Basis
+
+_basis_list = [b for g in instance.galaxies for b in g.cls_list_from(cls=_Basis)]
+n_linear_gaussians = sum(len(b.profile_list) for b in _basis_list)
+
+print("\n--- Configuration (determines run time) ---")
+print(f"  Instrument:              {instrument}")
+print(f"  Pixel scale:             {pixel_scale} arcsec/pixel")
+print(f"  Real-space mask radius:  {mask_radius} arcsec")
+print(f"  Real-space grid shape:   {real_space_shape[0]} x {real_space_shape[1]}")
+print(f"  Visibilities:            {n_visibilities}")
+print(f"  Linear Gaussians:        {n_linear_gaussians}")
+
+# ---------------------------------------------------------------------------
+# 5. Full-pipeline reference (FitInterferometer) — eager baseline
+# ---------------------------------------------------------------------------
+
+print("\n--- Full FitInterferometer (eager baseline) ---")
+
+with timer.section("fit_interferometer_eager"):
+    fit = al.FitInterferometer(
+        dataset=dataset,
+        tracer=tracer,
+        xp=np,
+    )
+    figure_of_merit_ref = fit.figure_of_merit
+    log_likelihood_ref = fit.log_likelihood
+
+print(f"  figure_of_merit = {figure_of_merit_ref}")
+print(f"  log_likelihood  = {log_likelihood_ref}")
+
+
+# ===================================================================
+# PART B — Full-pipeline JIT
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("FULL-PIPELINE JIT")
+print("=" * 70)
+
+analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)
+
+
+def full_pipeline_from_params(params_tree):
+    """Full interferometer likelihood from a pytree-shaped ``ModelInstance``.
+
+    No flat-vector unpacking inside the trace — the instance crosses the JIT
+    boundary directly, with constants (redshifts, etc.) kept static via the
+    ``aux_data`` partition set up by ``autofit.jax.register_model``.
+    """
+    return analysis.log_likelihood_function(instance=params_tree)
+
+
+_, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree)
+full_pipeline_per_call = timer.records[-1][1] / 10
+
+print(f"  full log_likelihood = {full_result}")
+
+# ===================================================================
+# PART C — vmap + correctness
+# ===================================================================
+
+print("\n--- vmap batched evaluation ---")
+
+batch_size = 3
+
+parameters = jax.tree_util.tree_map(
+    lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)),
+    params_tree,
+)
+
+vmapped_full = jax.jit(jax.vmap(full_pipeline_from_params))
+
+with timer.section("vmap_first_call"):
+    result_vmap = vmapped_full(parameters)
+    block(result_vmap)
+
+n_vmap_repeats = 10
+with timer.section(f"vmap_steady_x{n_vmap_repeats}"):
+    for _ in range(n_vmap_repeats):
+        result_vmap = vmapped_full(parameters)
+        block(result_vmap)
+
+vmap_batch_time = timer.records[-1][1] / n_vmap_repeats
+vmap_per_call = vmap_batch_time / batch_size
+vmap_speedup = full_pipeline_per_call / vmap_per_call
+
+print(f"  batch results = {result_vmap}")
+print(f"  vmap batch of {batch_size}:   {vmap_batch_time:.6f} s")
+print(f"  vmap per call:         {vmap_per_call:.6f} s")
+print(f"  single JIT per call:   {full_pipeline_per_call:.6f} s")
+print(f"  vmap speedup:          {vmap_speedup:.1f}x faster per likelihood")
+
+# Correctness: full-pipeline JIT must match eager FitInterferometer.log_likelihood
+# (NOT figure_of_merit — the analysis.log_likelihood_function returns the
+# log-likelihood scalar directly, whereas FitInterferometer.figure_of_merit may
+# be a log-evidence when an inversion is present).
+np.testing.assert_allclose(
+    float(full_result),
+    float(log_likelihood_ref),
+    rtol=1e-4,
+    err_msg="interferometer/mge: JIT log_likelihood does not match eager FitInterferometer",
+)
+print("  Eager-vs-JIT correctness PASSED")
+
+np.testing.assert_allclose(
+    np.array(result_vmap),
+    float(full_result),
+    rtol=1e-4,
+    err_msg="interferometer/mge: JAX vmap likelihood mismatch",
+)
+print("  vmap-vs-single-JIT correctness PASSED")
+
+# ===================================================================
+# PART D — Static memory analysis
+# ===================================================================
+
+print("\n--- Static memory analysis ---")
+
+lowered_batched = vmapped_full.lower(parameters)
+compiled_batched = lowered_batched.compile()
+
+memory_analysis = compiled_batched.memory_analysis()
+print(f"  Output size:  {memory_analysis.output_size_in_bytes / 1024**2:.3f} MB")
+print(f"  Temp size:    {memory_analysis.temp_size_in_bytes / 1024**2:.3f} MB")
+print(
+    f"  Total:        "
+    f"{(memory_analysis.output_size_in_bytes + memory_analysis.temp_size_in_bytes) / 1024**2:.3f} MB"
+)
+
+# ===================================================================
+# JAX Likelihood Function Summary + artefacts
+# ===================================================================
+
+import json
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+al_version = al.__version__
+
+print("\n" + "=" * 70)
+print(f"JAX LIKELIHOOD FUNCTION SUMMARY — {instrument.upper()} — v{al_version}")
+print("=" * 70)
+print(f"  Instrument:              {instrument}")
+print(f"  Pixel scale:             {pixel_scale} arcsec/pixel")
+print(f"  Real-space mask radius:  {mask_radius} arcsec")
+print(f"  Real-space grid shape:   {real_space_shape[0]} x {real_space_shape[1]}")
+print(f"  Visibilities:            {n_visibilities}")
+print(f"  Linear Gaussians:        {n_linear_gaussians}")
+print("-" * 70)
+print(f"  Eager log_likelihood:    {log_likelihood_ref}")
+print(f"  JIT  log_likelihood:     {float(full_result)}")
+print("-" * 70)
+print(f"  Full pipeline per call:  {full_pipeline_per_call:.6f} s")
+print(f"  vmap batch={batch_size} per call:   {vmap_per_call:.6f} s")
+print(f"  vmap speedup:            {vmap_speedup:.1f}x")
+print("=" * 70)
+
+# --- Save results dictionary ---
+
+likelihood_summary = {
+    "autolens_version": al_version,
+    "instrument": instrument,
+    "model": "mge",
+    "configuration": {
+        "pixel_scale_arcsec": pixel_scale,
+        "mask_radius_arcsec": mask_radius,
+        "real_space_shape": list(real_space_shape),
+        "visibilities": int(n_visibilities),
+        "linear_gaussians": int(n_linear_gaussians),
+    },
+    "log_likelihood_eager": float(log_likelihood_ref),
+    "log_likelihood_jit": float(full_result),
+    "full_pipeline_single_jit": full_pipeline_per_call,
+    "vmap": {
+        "batch_size": batch_size,
+        "batch_time": vmap_batch_time,
+        "per_call": vmap_per_call,
+        "speedup_vs_single_jit": round(vmap_speedup, 1),
+    },
+    "memory_mb": {
+        "output": memory_analysis.output_size_in_bytes / 1024**2,
+        "temp": memory_analysis.temp_size_in_bytes / 1024**2,
+    },
+}
+
+results_dir = _workspace_root / "results" / "likelihood" / "interferometer"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+dict_path = results_dir / f"mge_likelihood_summary_{instrument}_v{al_version}.json"
+dict_path.write_text(json.dumps(likelihood_summary, indent=2))
+print(f"\n  Results dict saved to: {dict_path}")
+
+# --- Save bar chart ---
+
+labels = [
+    f"Full pipeline (single JIT)",
+    f"vmap batch={batch_size} (per call)",
+]
+times = [full_pipeline_per_call, vmap_per_call]
+
+fig, ax = plt.subplots(figsize=(10, 3.5))
+y_pos = range(len(labels))
+bars = ax.barh(
+    y_pos, times, color=["#4C72B0", "#55A868"], edgecolor="white", height=0.55
+)
+
+for bar, t in zip(bars, times):
+    ax.text(
+        bar.get_width() + max(times) * 0.01,
+        bar.get_y() + bar.get_height() / 2,
+        f"{t:.6f} s",
+        va="center",
+        fontsize=9,
+    )
+
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=10)
+ax.invert_yaxis()
+ax.set_xlabel("Time per call (s)", fontsize=11)
+fig.suptitle(
+    f"MGE Interferometer Likelihood — {instrument.upper()}",
+    fontsize=12,
+    fontweight="bold",
+)
+ax.set_title(
+    f'AutoLens v{al_version}  |  {pixel_scale}"/px  |  '
+    f"{real_space_shape[0]}x{real_space_shape[1]} real-space  |  "
+    f"{n_visibilities} visibilities  |  {n_linear_gaussians} Gaussians  |  "
+    f"vmap speedup: {vmap_speedup:.1f}x",
+    fontsize=9,
+)
+ax.margins(x=0.2)
+fig.tight_layout()
+
+chart_path = results_dir / f"mge_likelihood_summary_{instrument}_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to:    {chart_path}")
+
+
+# ===================================================================
+# Regression assertion — realistic-scale deterministic log-likelihood
+# ===================================================================
+#
+# Simulator truth parameters via GaussianPrior(mean=truth, sigma=small)
+# make the full-pipeline log-likelihood deterministic at the prior median.
+EXPECTED_LOG_LIKELIHOOD_SMA = -3153.8939746810656
+
+np.testing.assert_allclose(
+    log_likelihood_ref,
+    EXPECTED_LOG_LIKELIHOOD_SMA,
+    rtol=1e-4,
+    err_msg=(
+        f"interferometer/mge[{instrument}]: regression — eager log_likelihood drifted "
+        f"(got {log_likelihood_ref}, expected {EXPECTED_LOG_LIKELIHOOD_SMA})"
+    ),
+)
+print(
+    f"  Eager regression assertion PASSED: log_likelihood matches "
+    f"{EXPECTED_LOG_LIKELIHOOD_SMA:.6f}"
+)
+np.testing.assert_allclose(
+    float(full_result),
+    EXPECTED_LOG_LIKELIHOOD_SMA,
+    rtol=1e-4,
+    err_msg=f"interferometer/mge[{instrument}]: regression — full log_likelihood drifted",
+)
+np.testing.assert_allclose(
+    np.array(result_vmap),
+    EXPECTED_LOG_LIKELIHOOD_SMA,
+    rtol=1e-4,
+    err_msg=f"interferometer/mge[{instrument}]: regression — vmap log_likelihood drifted",
+)
+print(
+    f"  Regression assertion PASSED: log_likelihood matches {EXPECTED_LOG_LIKELIHOOD_SMA:.6f}"
+)
diff --git a/likelihood/interferometer/pixelization.py b/likelihood/interferometer/pixelization.py
new file mode 100644
index 0000000..384df58
--- /dev/null
+++ b/likelihood/interferometer/pixelization.py
@@ -0,0 +1,558 @@
+"""
+JAX Profiling: Pixelization Interferometer Likelihood
+=====================================================
+
+Profiles the JAX likelihood function for an interferometer dataset where the
+source galaxy is reconstructed using a rectangular pixelization with constant
+regularization, and the lens galaxy is an Isothermal + ExternalShear.
+
+Mirrors ``likelihood/interferometer/mge.py`` (Phase 1) with the MGE source
+replaced by a ``RectangularUniform`` pixelization + ``Constant`` regularization
+— matching ``likelihood/imaging/pixelization.py`` so imaging vs
+interferometer results can be compared side-by-side.
+
+Like the MGE interferometer script, this script deliberately does not per-step
+JIT the inversion pipeline. The interferometer path exercises a
+Fourier-transformed mapping matrix, a visibilities-space data vector /
+curvature matrix, and an NNLS solve whose ``xp=jnp`` threading has not been
+fully characterised. Per-step decomposition risks missing cross-step XLA
+fusion that matters in practice, and risks hitting library-level JAX blockers
+that we would want to raise as separate issues rather than work around here.
+Once the full-pipeline JIT is stable on interferometer, the per-step
+breakdown can land as a follow-up.
+
+Measures:
+
+1. Eager baseline: ``FitInterferometer`` with ``xp=np``, print
+   ``figure_of_merit`` / ``log_likelihood``.
+2. Full-pipeline JIT: ``jax.jit(analysis.log_likelihood_function)`` on a
+   pytree-registered ``ModelInstance``. Measure lower / compile / first-call /
+   steady-state per-call.
+3. Batched evaluation: ``jax.jit(jax.vmap(full_pipeline))``. Measure
+   per-likelihood cost and speedup vs the single-JIT path.
+4. Correctness: eager vs JIT log-likelihood agreement at ``rtol=1e-4``.
+5. Static memory analysis of the batched program.
+6. Results JSON + PNG written to ``results/`` using the same schema as the
+   imaging profiling scripts so they can be compared side-by-side.
+
+Pytree-native parameter inputs
+------------------------------
+
+Uses ``af.ModelInstance`` as the JIT input via PyAutoFit's opt-in pytree
+registration (``autofit.jax.register_model``). Exercises the ``TuplePrior``
+pytree support landed in PyAutoFit#1222.
+"""
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import time
+import subprocess
+import sys
+from pathlib import Path
+from contextlib import contextmanager
+
+import autofit as af
+import autolens as al
+from autofit.jax import register_model as _register_model_pytrees
+
+# ---------------------------------------------------------------------------
+# Instrument configuration
+# ---------------------------------------------------------------------------
+
+INSTRUMENTS = {
+    "sma": {"pixel_scale": 0.1, "real_space_shape": (256, 256)},
+    "alma": {"pixel_scale": 0.05, "real_space_shape": (256, 256)},
+}
+
+instrument = "sma"  # <-- change this to profile a different instrument
+
+mesh_pixels_yx = 28
+mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)
+regularization_coefficient = 1.0
+
+
+# ---------------------------------------------------------------------------
+# Profiling helpers
+# ---------------------------------------------------------------------------
+
+
+class Timer:
+    """Accumulates named timing measurements and prints a summary."""
+
+    def __init__(self):
+        self.records: list[tuple[str, float]] = []
+
+    @contextmanager
+    def section(self, label: str):
+        start = time.perf_counter()
+        yield
+        elapsed = time.perf_counter() - start
+        self.records.append((label, elapsed))
+        print(f"  [{label}] {elapsed:.4f} s")
+
+
+def block(x):
+    """Call block_until_ready if available (JAX arrays)."""
+    if hasattr(x, "block_until_ready"):
+        x.block_until_ready()
+    return x
+
+
+def jit_profile(func, label, *args, n_repeats=10):
+    """JIT-compile *func*, time lower / compile / first call / steady state."""
+    jitted = jax.jit(func)
+
+    with timer.section(f"{label}_lower"):
+        lowered = jitted.lower(*args)
+
+    with timer.section(f"{label}_compile"):
+        compiled = lowered.compile()
+
+    with timer.section(f"{label}_first_call"):
+        result = compiled(*args)
+        block(result)
+
+    with timer.section(f"{label}_steady_x{n_repeats}"):
+        for _ in range(n_repeats):
+            result = compiled(*args)
+            block(result)
+
+    per_call = timer.records[-1][1] / n_repeats
+    print(f"    -> per-call avg: {per_call:.6f} s")
+    return compiled, result
+
+
+timer = Timer()
+
+# ===================================================================
+# PART A — Setup (not JIT-compiled)
+# ===================================================================
+
+# ---------------------------------------------------------------------------
+# 1. Dataset
+# ---------------------------------------------------------------------------
+
+print(f"\n--- Dataset loading [{instrument}] ---")
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[1]
+pixel_scale = INSTRUMENTS[instrument]["pixel_scale"]
+real_space_shape = INSTRUMENTS[instrument]["real_space_shape"]
+dataset_path = Path("dataset") / "interferometer" / instrument
+
+if al.util.dataset.should_simulate(str(dataset_path)):
+    raise FileNotFoundError(
+        f"Input dataset missing at '{dataset_path}'. The autolens_profiling "
+        f"repo mirrors only the curated datasets needed for default smoke "
+        f"runs. To regenerate or extend datasets, use the source-of-truth "
+        f"scripts under autolens_workspace_developer/jax_profiling/dataset_setup/, "
+        f"then copy the result into autolens_profiling/dataset/."
+    )
+
+mask_radius = 3.0
+
+real_space_mask = al.Mask2D.circular(
+    shape_native=real_space_shape,
+    pixel_scales=pixel_scale,
+    radius=mask_radius,
+)
+
+with timer.section("dataset_load"):
+    dataset = al.Interferometer.from_fits(
+        data_path=dataset_path / "data.fits",
+        noise_map_path=dataset_path / "noise_map.fits",
+        uv_wavelengths_path=dataset_path / "uv_wavelengths.fits",
+        real_space_mask=real_space_mask,
+        transformer_class=al.TransformerDFT,
+    )
+
+n_visibilities = dataset.uv_wavelengths.shape[0]
+print(f"  Total visibilities: {n_visibilities}")
+
+# ---------------------------------------------------------------------------
+# 2. Model construction
+# ---------------------------------------------------------------------------
+
+print("\n--- Model construction ---")
+
+with timer.section("model_build"):
+    # GaussianPrior(mean=truth, sigma=small) centres prior-median at the
+    # simulator truth while keeping params free so gradient diagnostics
+    # have dimensionality.
+    mass = af.Model(al.mp.Isothermal)
+    mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.005)
+    mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.005)
+    mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05)
+    _lens_mass_ell = al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0)
+    mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=_lens_mass_ell[0], sigma=0.01)
+    mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=_lens_mass_ell[1], sigma=0.01)
+
+    shear = af.Model(al.mp.ExternalShear)
+    shear.gamma_1 = af.GaussianPrior(mean=0.05, sigma=0.005)
+    shear.gamma_2 = af.GaussianPrior(mean=0.05, sigma=0.005)
+
+    lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)
+
+    pixelization = af.Model(
+        al.Pixelization,
+        mesh=al.mesh.RectangularUniform(shape=mesh_shape),
+        regularization=al.reg.Constant(coefficient=regularization_coefficient),
+    )
+
+    source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)
+
+    model = af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+print(f"  Total free parameters: {model.total_free_parameters}")
+
+# ---------------------------------------------------------------------------
+# 3. Instantiate concrete objects from prior medians
+# ---------------------------------------------------------------------------
+
+print("\n--- Instantiate concrete model ---")
+
+with timer.section("instance_from_vector"):
+    param_vector = model.physical_values_from_prior_medians
+    instance = model.instance_from_vector(vector=param_vector)
+
+with timer.section("register_pytrees"):
+    _register_model_pytrees(model)
+
+# JIT input: the instance itself, with all parameter leaves promoted to JAX
+# arrays. The eager NumPy instance is retained for the eager FitInterferometer
+# baseline below.
+params_tree = jax.tree_util.tree_map(jnp.asarray, instance)
+
+tracer = al.Tracer(galaxies=list(instance.galaxies))
+
+print(f"  Tracer planes: {tracer.total_planes}")
+
+# ---------------------------------------------------------------------------
+# 4. Configuration summary
+# ---------------------------------------------------------------------------
+
+n_source_pixels = mesh_shape[0] * mesh_shape[1]
+
+print("\n--- Configuration (determines run time) ---")
+print(f"  Instrument:              {instrument}")
+print(f"  Pixel scale:             {pixel_scale} arcsec/pixel")
+print(f"  Real-space mask radius:  {mask_radius} arcsec")
+print(f"  Real-space grid shape:   {real_space_shape[0]} x {real_space_shape[1]}")
+print(f"  Visibilities:            {n_visibilities}")
+print(f"  Mesh shape:              {mesh_shape[0]} x {mesh_shape[1]}")
+print(f"  Source pixels:           {n_source_pixels}")
+print(f"  Reg. coefficient:        {regularization_coefficient}")
+
+# ---------------------------------------------------------------------------
+# 5. Full-pipeline reference (FitInterferometer) — eager baseline
+# ---------------------------------------------------------------------------
+
+print("\n--- Full FitInterferometer (eager baseline) ---")
+
+with timer.section("fit_interferometer_eager"):
+    fit = al.FitInterferometer(
+        dataset=dataset,
+        tracer=tracer,
+        xp=np,
+    )
+    figure_of_merit_ref = fit.figure_of_merit
+    log_likelihood_ref = fit.log_likelihood
+
+print(f"  figure_of_merit = {figure_of_merit_ref}")
+print(f"  log_likelihood  = {log_likelihood_ref}")
+
+
+# ===================================================================
+# PART B — Full-pipeline JIT
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("FULL-PIPELINE JIT")
+print("=" * 70)
+
+analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)
+
+
+def full_pipeline_from_params(params_tree):
+    """Full interferometer likelihood from a pytree-shaped ``ModelInstance``.
+
+    No flat-vector unpacking inside the trace — the instance crosses the JIT
+    boundary directly, with constants (redshifts, etc.) kept static via the
+    ``aux_data`` partition set up by ``autofit.jax.register_model``.
+    """
+    return analysis.log_likelihood_function(instance=params_tree)
+
+
+_, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree)
+full_pipeline_per_call = timer.records[-1][1] / 10
+
+print(f"  full log_likelihood = {full_result}")
+
+# ===================================================================
+# PART C — vmap + correctness
+# ===================================================================
+#
+# NOTE: vmap requires at least one JAX array leaf in the params_tree.
+# When model.total_free_parameters == 0 (all params fixed to truth), the
+# pytree has no array leaves and vmap cannot batch over it. Skip in that case.
+
+print("\n--- vmap batched evaluation ---")
+
+batch_size = 3
+vmap_batch_time = None
+vmap_per_call = None
+vmap_speedup = None
+result_vmap = None
+
+_n_leaves = len(jax.tree_util.tree_leaves(params_tree))
+if _n_leaves == 0:
+    print(
+        f"  SKIPPED: model has 0 free parameters (all fixed to truth); "
+        f"vmap requires at least one array leaf."
+    )
+else:
+    parameters = jax.tree_util.tree_map(
+        lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)),
+        params_tree,
+    )
+
+    vmapped_full = jax.jit(jax.vmap(full_pipeline_from_params))
+
+    with timer.section("vmap_first_call"):
+        result_vmap = vmapped_full(parameters)
+        block(result_vmap)
+
+    n_vmap_repeats = 10
+    with timer.section(f"vmap_steady_x{n_vmap_repeats}"):
+        for _ in range(n_vmap_repeats):
+            result_vmap = vmapped_full(parameters)
+            block(result_vmap)
+
+    vmap_batch_time = timer.records[-1][1] / n_vmap_repeats
+    vmap_per_call = vmap_batch_time / batch_size
+    vmap_speedup = full_pipeline_per_call / vmap_per_call
+
+    print(f"  batch results = {result_vmap}")
+    print(f"  vmap batch of {batch_size}:   {vmap_batch_time:.6f} s")
+    print(f"  vmap per call:         {vmap_per_call:.6f} s")
+    print(f"  single JIT per call:   {full_pipeline_per_call:.6f} s")
+    print(f"  vmap speedup:          {vmap_speedup:.1f}x faster per likelihood")
+
+    # Correctness: for inversion models (pixelization + regularization), the
+    # analysis "log_likelihood_function" actually returns the log-evidence
+    # (= figure_of_merit), which includes the regularization/determinant terms.
+    # Match against figure_of_merit_ref, not log_likelihood_ref. For the MGE case
+    # in Phase 1 the two were numerically equal because there was no regularized
+    # inversion; here they differ by the regularization + log-determinant terms.
+    np.testing.assert_allclose(
+        float(full_result),
+        float(figure_of_merit_ref),
+        rtol=1e-4,
+        err_msg="interferometer/pixelization: JIT log-evidence does not match eager figure_of_merit",
+    )
+    print("  Eager-vs-JIT correctness PASSED")
+
+    np.testing.assert_allclose(
+        np.array(result_vmap),
+        float(full_result),
+        rtol=1e-4,
+        err_msg="interferometer/pixelization: JAX vmap likelihood mismatch",
+    )
+    print("  vmap-vs-single-JIT correctness PASSED")
+
+# ===================================================================
+# PART D — Static memory analysis
+# ===================================================================
+
+print("\n--- Static memory analysis ---")
+
+if _n_leaves == 0:
+    print("  SKIPPED: no array leaves in params_tree (all params fixed to truth).")
+    memory_analysis = None
+else:
+    lowered_batched = vmapped_full.lower(parameters)
+    compiled_batched = lowered_batched.compile()
+
+    memory_analysis = compiled_batched.memory_analysis()
+    print(f"  Output size:  {memory_analysis.output_size_in_bytes / 1024**2:.3f} MB")
+    print(f"  Temp size:    {memory_analysis.temp_size_in_bytes / 1024**2:.3f} MB")
+    print(
+        f"  Total:        "
+        f"{(memory_analysis.output_size_in_bytes + memory_analysis.temp_size_in_bytes) / 1024**2:.3f} MB"
+    )
+
+# ===================================================================
+# JAX Likelihood Function Summary + artefacts
+# ===================================================================
+
+import json
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+al_version = al.__version__
+
+print("\n" + "=" * 70)
+print(f"JAX LIKELIHOOD FUNCTION SUMMARY — {instrument.upper()} — v{al_version}")
+print("=" * 70)
+print(f"  Instrument:              {instrument}")
+print(f"  Pixel scale:             {pixel_scale} arcsec/pixel")
+print(f"  Real-space mask radius:  {mask_radius} arcsec")
+print(f"  Real-space grid shape:   {real_space_shape[0]} x {real_space_shape[1]}")
+print(f"  Visibilities:            {n_visibilities}")
+print(f"  Mesh shape:              {mesh_shape[0]} x {mesh_shape[1]}")
+print(f"  Source pixels:           {n_source_pixels}")
+print("-" * 70)
+print(f"  Eager log_likelihood:    {log_likelihood_ref}")
+print(f"  Eager figure_of_merit:   {figure_of_merit_ref}  (log-evidence)")
+print(f"  JIT  log-evidence:       {float(full_result)}")
+print("-" * 70)
+print(f"  Full pipeline per call:  {full_pipeline_per_call:.6f} s")
+if vmap_per_call is not None:
+    print(f"  vmap batch={batch_size} per call:   {vmap_per_call:.6f} s")
+    print(f"  vmap speedup:            {vmap_speedup:.1f}x")
+else:
+    print(f"  vmap:                    SKIPPED (0 free params)")
+print("=" * 70)
+
+# --- Save results dictionary ---
+
+likelihood_summary = {
+    "autolens_version": al_version,
+    "instrument": instrument,
+    "model": "pixelization",
+    "configuration": {
+        "pixel_scale_arcsec": pixel_scale,
+        "mask_radius_arcsec": mask_radius,
+        "real_space_shape": list(real_space_shape),
+        "visibilities": int(n_visibilities),
+        "mesh_shape": list(mesh_shape),
+        "source_pixels": int(n_source_pixels),
+        "regularization_coefficient": regularization_coefficient,
+    },
+    "log_likelihood_eager": float(log_likelihood_ref),
+    "figure_of_merit_eager": float(figure_of_merit_ref),
+    "log_evidence_jit": float(full_result),
+    "full_pipeline_single_jit": full_pipeline_per_call,
+    "vmap": (
+        "SKIPPED — model has 0 free parameters (all fixed to truth)"
+        if vmap_per_call is None
+        else {
+            "batch_size": batch_size,
+            "batch_time": vmap_batch_time,
+            "per_call": vmap_per_call,
+            "speedup_vs_single_jit": round(vmap_speedup, 1),
+        }
+    ),
+    "memory_mb": (
+        None
+        if memory_analysis is None
+        else {
+            "output": memory_analysis.output_size_in_bytes / 1024**2,
+            "temp": memory_analysis.temp_size_in_bytes / 1024**2,
+        }
+    ),
+}
+
+results_dir = _workspace_root / "results" / "likelihood" / "interferometer"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+dict_path = (
+    results_dir / f"pixelization_likelihood_summary_{instrument}_v{al_version}.json"
+)
+dict_path.write_text(json.dumps(likelihood_summary, indent=2))
+print(f"\n  Results dict saved to: {dict_path}")
+
+# --- Save bar chart ---
+
+labels = [f"Full pipeline (single JIT)"]
+times = [full_pipeline_per_call]
+bar_colors = ["#4C72B0"]
+if vmap_per_call is not None:
+    labels.append(f"vmap batch={batch_size} (per call)")
+    times.append(vmap_per_call)
+    bar_colors.append("#55A868")
+
+fig, ax = plt.subplots(figsize=(10, 3.5))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color=bar_colors, edgecolor="white", height=0.55)
+
+for bar, t in zip(bars, times):
+    ax.text(
+        bar.get_width() + max(times) * 0.01,
+        bar.get_y() + bar.get_height() / 2,
+        f"{t:.6f} s",
+        va="center",
+        fontsize=9,
+    )
+
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=10)
+ax.invert_yaxis()
+ax.set_xlabel("Time per call (s)", fontsize=11)
+fig.suptitle(
+    f"Pixelization Interferometer Likelihood — {instrument.upper()}",
+    fontsize=12,
+    fontweight="bold",
+)
+_vmap_title = (
+    f"vmap speedup: {vmap_speedup:.1f}x"
+    if vmap_speedup is not None
+    else "vmap: SKIPPED"
+)
+ax.set_title(
+    f'AutoLens v{al_version}  |  {pixel_scale}"/px  |  '
+    f"{real_space_shape[0]}x{real_space_shape[1]} real-space  |  "
+    f"{n_visibilities} visibilities  |  {mesh_shape[0]}x{mesh_shape[1]} mesh  |  "
+    f"{_vmap_title}",
+    fontsize=9,
+)
+ax.margins(x=0.2)
+fig.tight_layout()
+
+chart_path = (
+    results_dir / f"pixelization_likelihood_summary_{instrument}_v{al_version}.png"
+)
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to:    {chart_path}")
+
+
+# ===================================================================
+# Regression assertion — realistic-scale deterministic log-evidence
+# ===================================================================
+#
+# Simulator truth parameters via GaussianPrior(mean=truth, sigma=small)
+# make the full-pipeline log-evidence deterministic at the prior median.
+EXPECTED_LOG_EVIDENCE_SMA = -3165.251161569041
+
+np.testing.assert_allclose(
+    figure_of_merit_ref,
+    EXPECTED_LOG_EVIDENCE_SMA,
+    rtol=1e-4,
+    err_msg=(
+        f"interferometer/pixelization[{instrument}]: regression — eager log_evidence "
+        f"drifted (got {figure_of_merit_ref}, expected {EXPECTED_LOG_EVIDENCE_SMA})"
+    ),
+)
+print(
+    f"  Eager regression assertion PASSED: log_evidence matches "
+    f"{EXPECTED_LOG_EVIDENCE_SMA:.6f}"
+)
+np.testing.assert_allclose(
+    float(full_result),
+    EXPECTED_LOG_EVIDENCE_SMA,
+    rtol=1e-4,
+    err_msg=f"interferometer/pixelization[{instrument}]: regression — full log_evidence drifted",
+)
+np.testing.assert_allclose(
+    np.array(result_vmap),
+    EXPECTED_LOG_EVIDENCE_SMA,
+    rtol=1e-4,
+    err_msg=f"interferometer/pixelization[{instrument}]: regression — vmap log_evidence drifted",
+)
+print(
+    f"  Regression assertion PASSED: log_evidence matches {EXPECTED_LOG_EVIDENCE_SMA:.6f}"
+)
diff --git a/likelihood/point_source/README.md b/likelihood/point_source/README.md
new file mode 100644
index 0000000..63b6a76
--- /dev/null
+++ b/likelihood/point_source/README.md
@@ -0,0 +1,55 @@
+# likelihood/point_source
+
+JAX JIT profiling for the PyAutoLens **point-source** likelihood function — fitting the multiple images of a lensed point source (quasar, supernova, or compact emitter) given an observed set of image-plane positions and a parametric lens model.
+
+The point-source likelihood has two distinct chi-squared variants depending on where the comparison between model and data happens. Both are JIT-traceable end-to-end through the underlying `PointSolver` (a triangle-refinement loop that threads `xp=jnp` through every step).
+
+## Two chi-squared variants
+
+| Variant | Computed in | Comparison |
+|---------|-------------|------------|
+| **Image-plane** | image plane coordinates | Solve for the model image positions via `PointSolver`, pair each model image with the closest observed image, compute χ² in image-plane coordinates. |
+| **Source-plane** | source plane coordinates | Ray-trace observed image positions back to the source plane, compute χ² of the inferred source-plane scatter (i.e. how tightly the back-traced images cluster around a single source). |
+
+Image-plane fitting is closer to what an observer would naively expect ("how well do the model images line up with the data images?"), but is more expensive because of the forward solver. Source-plane fitting is cheaper (no forward solve) but is a proxy: zero scatter in the source plane does not exactly correspond to zero image-plane chi-squared.
+
+## What each script measures
+
+For both variants:
+
+1. **Eager baseline** — `xp=np` log-likelihood for sanity.
+2. **Single-JIT** — lower / compile / first-call / steady-state per-call timings.
+3. **vmap** — batched per-likelihood cost and speedup vs single-JIT.
+4. **Three-tier numerical correctness**:
+    - eager ≡ JIT
+    - JIT ≡ vmap (every entry of the batched output)
+    - a hardcoded `EXPECTED_LOG_LIKELIHOOD_*` regression constant guarding against silent drift in the solver / chi-squared stack. This depends on the seeded simulator (`noise_seed=1` in the upstream `dataset_setup/point_source.py`) staying bit-stable.
+
+## Scripts
+
+| Script | Variant | Notes |
+|--------|---------|-------|
+| [`image_plane.py`](./image_plane.py) | Image-plane χ² | Full pipeline JIT-traceable end-to-end because `PointSolver` threads `xp=jnp` through every step. |
+| [`source_plane.py`](./source_plane.py) | Source-plane χ² | Cheaper than image-plane; no forward solver. |
+
+## Default dataset
+
+`dataset/point_source/simple/` — a minimal seeded dataset with `point_dataset_positions_only.json` (4 observed image positions) and the truth `tracer.json`. Both files are committed to this repo.
+
+## Headline run-times (populated by Phase 4)
+
+| Script | Dataset | CPU | Laptop GPU | A100 |
+|--------|---------|-----|------------|------|
+| `image_plane.py` | simple | _populated_ | _populated_ | _populated_ |
+| `source_plane.py` | simple | _populated_ | _populated_ | _populated_ |
+
+Numbers are the **steady-state per-call cost** (single-JIT, post-warmup), in milliseconds. Phase 4's dashboard auto-fills this from the latest `*_summary_v<version>.json` artifacts under `results/likelihood/point_source/`.
+
+## Output
+
+Each script writes:
+
+```
+results/likelihood/point_source/<script>_likelihood_summary_<dataset_name>_v<al.__version__>.json
+results/likelihood/point_source/<script>_likelihood_summary_<dataset_name>_v<al.__version__>.png
+```
diff --git a/likelihood/point_source/image_plane.py b/likelihood/point_source/image_plane.py
new file mode 100644
index 0000000..b8278fc
--- /dev/null
+++ b/likelihood/point_source/image_plane.py
@@ -0,0 +1,476 @@
+"""
+JAX Profiling: Point-Source Likelihood — Image-Plane Chi-Squared
+=================================================================
+
+Profiles ``AnalysisPoint.log_likelihood_function`` for a lensed point-source
+``PointDataset`` using the **image-plane** chi-squared
+(``al.FitPositionsImagePairAll``).
+
+Image-plane fitting solves for the model multiple-image positions in the
+image plane via the ``PointSolver`` (which JIT-traces a triangle-refinement
+loop), pairs each model image with the closest observed image, and computes
+a chi-squared in image-plane coordinates.
+
+Unlike the source-plane variant (see ``source_plane.py``), the full
+image-plane pipeline IS JIT-traceable end-to-end because ``PointSolver``
+threads ``xp=jnp`` through every step and ``FitPositionsImagePairAll``
+constructs its model-data via JAX-friendly operations.
+
+Pytree-native parameter inputs
+------------------------------
+
+This script uses ``af.ModelInstance`` as the JIT input via PyAutoFit's
+opt-in pytree registration (``autofit.jax.register_model``, PRs #1220 /
+#1221 / #1222).  The JIT'd closure consumes the registered instance
+directly, mirroring the pattern in ``../imaging/mge.py``.
+
+Three-tier numerical assertions
+-------------------------------
+
+1. **eager ≡ JIT**: numpy-path log-likelihood matches single-JIT result.
+2. **JIT ≡ vmap**: every entry of the batched vmap output matches the
+   single-JIT result.
+3. **regression constant**: hardcoded
+   ``EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE`` guards against silent drift in
+   the underlying solver / chi-squared stack.  This depends on the seeded
+   simulator (``noise_seed=1`` in ``simulators/point_source.py``) staying
+   bit-stable.
+"""
+
+import json
+import time
+import subprocess
+import sys
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+import autofit as af
+import autolens as al
+from autofit.jax import register_model as _register_model_pytrees
+
+
+# ---------------------------------------------------------------------------
+# Profiling helpers (mirrors imaging/mge.py and source_plane.py)
+# ---------------------------------------------------------------------------
+
+
+class Timer:
+    def __init__(self):
+        self.records: list[tuple[str, float]] = []
+
+    @contextmanager
+    def section(self, label: str):
+        start = time.perf_counter()
+        yield
+        elapsed = time.perf_counter() - start
+        self.records.append((label, elapsed))
+        print(f"  [{label}] {elapsed:.4f} s")
+
+    def summary(self):
+        print("\n" + "=" * 70)
+        print("PROFILING SUMMARY")
+        print("=" * 70)
+        max_label = max(len(r[0]) for r in self.records)
+        total = 0.0
+        for label, elapsed in self.records:
+            print(f"  {label:<{max_label}}  {elapsed:>10.4f} s")
+            total += elapsed
+        print("-" * 70)
+        print(f"  {'TOTAL':<{max_label}}  {total:>10.4f} s")
+        print("=" * 70)
+
+
+def block(x):
+    if hasattr(x, "block_until_ready"):
+        x.block_until_ready()
+    return x
+
+
+def jit_profile(func, label, *args, n_repeats=10):
+    jitted = jax.jit(func)
+
+    with timer.section(f"{label}_lower"):
+        lowered = jitted.lower(*args)
+
+    with timer.section(f"{label}_compile"):
+        compiled = lowered.compile()
+
+    with timer.section(f"{label}_first_call"):
+        result = compiled(*args)
+        block(result)
+
+    with timer.section(f"{label}_steady_x{n_repeats}"):
+        for _ in range(n_repeats):
+            result = compiled(*args)
+            block(result)
+
+    per_call = timer.records[-1][1] / n_repeats
+    print(f"    -> per-call avg: {per_call:.6f} s")
+    return compiled, result
+
+
+timer = Timer()
+dataset_name = "simple"
+
+
+# ===================================================================
+# PART A — Setup
+# ===================================================================
+
+print(f"\n--- Dataset loading [{dataset_name}] ---")
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[1]
+dataset_path = Path("dataset") / "point_source" / dataset_name
+
+if al.util.dataset.should_simulate(str(dataset_path)):
+    raise FileNotFoundError(
+        f"Input dataset missing at '{dataset_path}'. The autolens_profiling "
+        f"repo mirrors only the curated datasets needed for default smoke "
+        f"runs. To regenerate or extend datasets, use the source-of-truth "
+        f"scripts under autolens_workspace_developer/jax_profiling/dataset_setup/, "
+        f"then copy the result into autolens_profiling/dataset/."
+    )
+
+with timer.section("dataset_load"):
+    dataset = al.from_json(
+        file_path=dataset_path / "point_dataset_positions_only.json",
+    )
+
+n_observed_positions = dataset.positions.shape[0]
+positions_noise_sigma = float(dataset.positions_noise_map[0])
+
+print("\n--- Point solver ---")
+
+with timer.section("solver_build"):
+    grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.2)
+    solver = al.PointSolver.for_grid(
+        grid=grid,
+        pixel_scale_precision=0.001,
+        magnification_threshold=0.1,
+    )
+
+print("\n--- Model construction ---")
+
+with timer.section("model_build"):
+    # GaussianPrior(mean=truth, sigma=small) centres prior-median at the
+    # simulator truth while keeping params free so gradient diagnostics
+    # have dimensionality.
+    mass = af.Model(al.mp.Isothermal)
+    mass.centre.centre_0 = af.GaussianPrior(mean=0.01, sigma=0.005)
+    mass.centre.centre_1 = af.GaussianPrior(mean=0.01, sigma=0.005)
+    mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05)
+    mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.01, sigma=0.01)
+    mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.01, sigma=0.01)
+    lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)
+
+    point_0 = af.Model(al.ps.PointFlux)
+    point_0.centre.centre_0 = af.GaussianPrior(mean=0.07, sigma=0.005)
+    point_0.centre.centre_1 = af.GaussianPrior(mean=0.07, sigma=0.005)
+    source = af.Model(al.Galaxy, redshift=1.0, point_0=point_0)
+
+    model = af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+print(f"  Total free parameters: {model.total_free_parameters}")
+
+print("\n--- Instantiate concrete model ---")
+
+with timer.section("instance_from_vector"):
+    param_vector = model.physical_values_from_prior_medians
+    instance = model.instance_from_vector(vector=param_vector)
+
+with timer.section("register_pytrees"):
+    _register_model_pytrees(model)
+
+params_tree = jax.tree_util.tree_map(jnp.asarray, instance)
+
+
+# ---------------------------------------------------------------------------
+# Eager baseline — full FitPointDataset (image-plane chi-squared)
+# ---------------------------------------------------------------------------
+
+print("\n--- Eager FitPointDataset (image-plane) ---")
+
+analysis_eager = al.AnalysisPoint(
+    dataset=dataset,
+    solver=solver,
+    fit_positions_cls=al.FitPositionsImagePairAll,
+    use_jax=False,
+)
+
+with timer.section("fit_eager"):
+    fit_eager = analysis_eager.fit_from(instance=instance)
+    log_likelihood_ref = float(fit_eager.log_likelihood)
+    figure_of_merit_ref = float(fit_eager.figure_of_merit)
+
+n_eager_repeats = 10
+with timer.section(f"eager_log_likelihood_x{n_eager_repeats}"):
+    for _ in range(n_eager_repeats):
+        analysis_eager.log_likelihood_function(instance=instance)
+eager_per_call = timer.records[-1][1] / n_eager_repeats
+
+print(f"  log_likelihood   = {log_likelihood_ref}")
+print(f"  figure_of_merit  = {figure_of_merit_ref}")
+print(f"  eager per-call   = {eager_per_call:.6f} s")
+
+
+# ===================================================================
+# PART B — Full-pipeline JIT
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("FULL-PIPELINE JIT (image-plane)")
+print("=" * 70)
+
+analysis_jax = al.AnalysisPoint(
+    dataset=dataset,
+    solver=solver,
+    fit_positions_cls=al.FitPositionsImagePairAll,
+    use_jax=True,
+)
+
+
+def full_pipeline_from_params(params_tree):
+    return analysis_jax.log_likelihood_function(instance=params_tree)
+
+
+_, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree)
+full_pipeline_per_call = timer.records[-1][1] / 10
+print(f"  full log_likelihood = {full_result}")
+
+
+# ===================================================================
+# PART C — vmap over the full pipeline
+# ===================================================================
+
+print("\n--- vmap batched evaluation ---")
+
+batch_size = 3
+
+batched_params = jax.tree_util.tree_map(
+    lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)),
+    params_tree,
+)
+
+vmapped_full = jax.jit(jax.vmap(full_pipeline_from_params))
+
+with timer.section("vmap_first_call"):
+    result_vmap = vmapped_full(batched_params)
+    block(result_vmap)
+
+n_vmap_repeats = 10
+with timer.section(f"vmap_steady_x{n_vmap_repeats}"):
+    for _ in range(n_vmap_repeats):
+        result_vmap = vmapped_full(batched_params)
+        block(result_vmap)
+
+vmap_batch_time = timer.records[-1][1] / n_vmap_repeats
+vmap_per_call = vmap_batch_time / batch_size
+vmap_speedup = full_pipeline_per_call / vmap_per_call
+
+print(f"  batch results = {result_vmap}")
+print(f"  vmap batch of {batch_size}:   {vmap_batch_time:.6f} s")
+print(f"  vmap per call:         {vmap_per_call:.6f} s")
+print(f"  single JIT per call:   {full_pipeline_per_call:.6f} s")
+print(f"  vmap speedup:          {vmap_speedup:.1f}x faster per likelihood")
+
+
+# ===================================================================
+# PART D — Three-tier numerical assertions
+# ===================================================================
+#
+# Tier 1: eager (NumPy path) ≡ single JIT
+# Tier 2: single JIT ≡ every entry of vmap output
+# Tier 3: hardcoded regression constant (deterministic via seeded simulator)
+
+np.testing.assert_allclose(
+    log_likelihood_ref,
+    float(full_result),
+    rtol=1e-4,
+    err_msg=(
+        f"point_source/image_plane: eager vs JIT mismatch — "
+        f"eager={log_likelihood_ref} vs JIT={float(full_result)}"
+    ),
+)
+print("  Tier 1: eager ≡ JIT assertion PASSED")
+
+np.testing.assert_allclose(
+    np.array(result_vmap),
+    float(full_result),
+    rtol=1e-4,
+    err_msg="point_source/image_plane: JIT vs vmap mismatch",
+)
+print("  Tier 2: JIT ≡ vmap assertion PASSED")
+
+
+# ===================================================================
+# PART E — Static memory analysis
+# ===================================================================
+
+print("\n--- Static memory analysis ---")
+
+lowered_batched = vmapped_full.lower(batched_params)
+compiled_batched = lowered_batched.compile()
+mem = compiled_batched.memory_analysis()
+print(f"  Output size:  {mem.output_size_in_bytes / 1024**2:.3f} MB")
+print(f"  Temp size:    {mem.temp_size_in_bytes / 1024**2:.3f} MB")
+print(
+    f"  Total:        "
+    f"{(mem.output_size_in_bytes + mem.temp_size_in_bytes) / 1024**2:.3f} MB"
+)
+
+
+# ===================================================================
+# Summary + outputs
+# ===================================================================
+
+al_version = al.__version__
+
+print("\n" + "=" * 70)
+print(f"JAX LIKELIHOOD SUMMARY — POINT SOURCE IMAGE-PLANE — v{al_version}")
+print("=" * 70)
+print(f"  Dataset:                    {dataset_name}")
+print(f"  Observed image positions:   {n_observed_positions}")
+print(f"  Position noise sigma:       {positions_noise_sigma}")
+print(f"  Free parameters:            {model.total_free_parameters}")
+print(
+    f"  fit_positions_cls:          FitPositionsImagePairAll (image-plane chi-squared)"
+)
+print("-" * 70)
+print(
+    f"  Eager full likelihood:      {eager_per_call:.6f} s/call  ({log_likelihood_ref:.6f})"
+)
+print(f"  Full pipeline (JIT):        {full_pipeline_per_call:.6f} s/call")
+print(f"  vmap per-call (batch={batch_size}):    {vmap_per_call:.6f} s")
+print(f"  vmap speedup vs single JIT:           {vmap_speedup:.1f}x")
+print("=" * 70)
+
+likelihood_summary = {
+    "autolens_version": al_version,
+    "dataset": dataset_name,
+    "fit_positions_cls": "FitPositionsImagePairAll",
+    "configuration": {
+        "observed_image_positions": int(n_observed_positions),
+        "positions_noise_sigma": positions_noise_sigma,
+        "free_parameters": int(model.total_free_parameters),
+    },
+    "eager_per_call": eager_per_call,
+    "eager_log_likelihood": log_likelihood_ref,
+    "full_pipeline_single_jit": full_pipeline_per_call,
+    "full_pipeline_log_likelihood": float(full_result),
+    "vmap": {
+        "batch_size": batch_size,
+        "batch_time": vmap_batch_time,
+        "per_call": vmap_per_call,
+        "speedup_vs_single_jit": round(vmap_speedup, 1),
+    },
+}
+
+results_dir = _workspace_root / "results" / "likelihood" / "point_source"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+dict_path = results_dir / f"image_plane_summary_v{al_version}.json"
+dict_path.write_text(json.dumps(likelihood_summary, indent=2))
+print(f"\n  Results dict saved to: {dict_path}")
+
+# --- Bar chart ---
+
+labels = [
+    "Eager full likelihood",
+    "Full pipeline (JIT)",
+    f"vmap per-call (batch={batch_size})",
+]
+times = [eager_per_call, full_pipeline_per_call, vmap_per_call]
+colors = ["#8172B3", "#C44E52", "#55A868"]
+
+fig, ax = plt.subplots(figsize=(10, 4.0))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color=colors, edgecolor="white", height=0.6)
+
+for bar, t in zip(bars, times):
+    ax.text(
+        bar.get_width() + max(times) * 0.01,
+        bar.get_y() + bar.get_height() / 2,
+        f"{t:.6f} s",
+        va="center",
+        fontsize=9,
+    )
+
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=10)
+ax.invert_yaxis()
+ax.set_xlabel("Time per call (s)", fontsize=11)
+fig.suptitle(
+    "Point-Source Likelihood — Image-Plane Chi-Squared",
+    fontsize=12,
+    fontweight="bold",
+)
+ax.set_title(
+    f"AutoLens v{al_version}  |  {n_observed_positions} positions  |  "
+    f"{model.total_free_parameters} free params  |  "
+    f"vmap speedup: {vmap_speedup:.1f}x",
+    fontsize=9,
+)
+ax.margins(x=0.20)
+fig.tight_layout()
+
+chart_path = results_dir / f"image_plane_summary_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to:    {chart_path}")
+
+
+# ===================================================================
+# Tier 3: regression assertion — deterministic via seeded simulator
+# ===================================================================
+#
+# Simulator truth parameters + seeded noise (noise_seed=1 in
+# simulators/point_source.py) make the image-plane log-likelihood
+# deterministic. Eager, JIT, and vmap all agree to float64.
+EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE = 0.07475703623045682
+
+np.testing.assert_allclose(
+    log_likelihood_ref,
+    EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE,
+    rtol=1e-4,
+    err_msg=(
+        f"point_source/image_plane: regression — eager log_likelihood drifted "
+        f"(got {log_likelihood_ref}, expected {EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE})"
+    ),
+)
+print(
+    f"  Eager regression assertion PASSED: log_likelihood matches "
+    f"{EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE:.6f}"
+)
+np.testing.assert_allclose(
+    float(full_result),
+    EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE,
+    rtol=1e-4,
+    err_msg=(
+        f"point_source/image_plane: regression — JIT log_likelihood drifted "
+        f"(got {float(full_result)}, expected {EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE})"
+    ),
+)
+np.testing.assert_allclose(
+    np.array(result_vmap),
+    EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE,
+    rtol=1e-4,
+    err_msg=(
+        f"point_source/image_plane: regression — vmap log_likelihood drifted "
+        f"(expected {EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE})"
+    ),
+)
+print(
+    f"  Tier 3: regression assertion PASSED: log_likelihood matches "
+    f"{EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE:.6f}"
+)
+
+timer.summary()
diff --git a/likelihood/point_source/source_plane.py b/likelihood/point_source/source_plane.py
new file mode 100644
index 0000000..43ab685
--- /dev/null
+++ b/likelihood/point_source/source_plane.py
@@ -0,0 +1,517 @@
+"""
+JAX Profiling: Point-Source Likelihood — Source-Plane Chi-Squared
+==================================================================
+
+Profiles ``AnalysisPoint.log_likelihood_function`` for a lensed point-source
+``PointDataset`` using the **source-plane** chi-squared
+(``al.FitPositionsSource``).
+
+Source-plane fitting traces each *observed* image-plane position back to the
+source plane via the lens model, then computes a chi-squared between the
+ray-traced positions and the model source position.  No image-plane solver
+is required.
+
+Pytree-native parameter inputs
+------------------------------
+
+This script uses ``af.ModelInstance`` as the JIT input via PyAutoFit's opt-in
+pytree registration (``autofit.jax.register_model``, PR #1220 / #1221 / #1222).
+The JIT'd closure consumes the registered instance directly.
+"""
+
+import json
+import time
+import subprocess
+import sys
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+import autofit as af
+import autoarray as aa
+import autolens as al
+from autofit.jax import register_model as _register_model_pytrees
+
+
+# ---------------------------------------------------------------------------
+# Profiling helpers (mirrors imaging/mge.py)
+# ---------------------------------------------------------------------------
+
+
+class Timer:
+    def __init__(self):
+        self.records: list[tuple[str, float]] = []
+
+    @contextmanager
+    def section(self, label: str):
+        start = time.perf_counter()
+        yield
+        elapsed = time.perf_counter() - start
+        self.records.append((label, elapsed))
+        print(f"  [{label}] {elapsed:.4f} s")
+
+    def summary(self):
+        print("\n" + "=" * 70)
+        print("PROFILING SUMMARY")
+        print("=" * 70)
+        max_label = max(len(r[0]) for r in self.records)
+        total = 0.0
+        for label, elapsed in self.records:
+            print(f"  {label:<{max_label}}  {elapsed:>10.4f} s")
+            total += elapsed
+        print("-" * 70)
+        print(f"  {'TOTAL':<{max_label}}  {total:>10.4f} s")
+        print("=" * 70)
+
+
+def block(x):
+    if hasattr(x, "block_until_ready"):
+        x.block_until_ready()
+    return x
+
+
+def jit_profile(func, label, *args, n_repeats=10):
+    jitted = jax.jit(func)
+
+    with timer.section(f"{label}_lower"):
+        lowered = jitted.lower(*args)
+
+    with timer.section(f"{label}_compile"):
+        compiled = lowered.compile()
+
+    with timer.section(f"{label}_first_call"):
+        result = compiled(*args)
+        block(result)
+
+    with timer.section(f"{label}_steady_x{n_repeats}"):
+        for _ in range(n_repeats):
+            result = compiled(*args)
+            block(result)
+
+    per_call = timer.records[-1][1] / n_repeats
+    print(f"    -> per-call avg: {per_call:.6f} s")
+    return compiled, result
+
+
+timer = Timer()
+dataset_name = "simple"
+
+
+# ===================================================================
+# PART A — Setup
+# ===================================================================
+
+print(f"\n--- Dataset loading [{dataset_name}] ---")
+
+_script_dir = Path(__file__).resolve().parent
+_workspace_root = _script_dir.parents[1]
+dataset_path = Path("dataset") / "point_source" / dataset_name
+
+if al.util.dataset.should_simulate(str(dataset_path)):
+    raise FileNotFoundError(
+        f"Input dataset missing at '{dataset_path}'. The autolens_profiling "
+        f"repo mirrors only the curated datasets needed for default smoke "
+        f"runs. To regenerate or extend datasets, use the source-of-truth "
+        f"scripts under autolens_workspace_developer/jax_profiling/dataset_setup/, "
+        f"then copy the result into autolens_profiling/dataset/."
+    )
+
+with timer.section("dataset_load"):
+    dataset = al.from_json(
+        file_path=dataset_path / "point_dataset_positions_only.json",
+    )
+
+n_observed_positions = dataset.positions.shape[0]
+positions_noise_sigma = float(dataset.positions_noise_map[0])
+
+print("\n--- Point solver ---")
+
+with timer.section("solver_build"):
+    grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.2)
+    solver = al.PointSolver.for_grid(
+        grid=grid,
+        pixel_scale_precision=0.001,
+        magnification_threshold=0.1,
+    )
+
+print("\n--- Model construction ---")
+
+with timer.section("model_build"):
+    # GaussianPrior(mean=truth, sigma=small) centres prior-median at the
+    # simulator truth while keeping params free so gradient vectors and
+    # finite-difference diagnostics have dimensionality.
+    mass = af.Model(al.mp.Isothermal)
+    mass.centre.centre_0 = af.GaussianPrior(mean=0.01, sigma=0.005)
+    mass.centre.centre_1 = af.GaussianPrior(mean=0.01, sigma=0.005)
+    mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05)
+    mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.01, sigma=0.005)
+    mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.01, sigma=0.005)
+    lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)
+
+    point_0 = af.Model(al.ps.PointFlux)
+    point_0.centre.centre_0 = af.GaussianPrior(mean=0.07, sigma=0.005)
+    point_0.centre.centre_1 = af.GaussianPrior(mean=0.07, sigma=0.005)
+    source = af.Model(al.Galaxy, redshift=1.0, point_0=point_0)
+
+    model = af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+print(f"  Total free parameters: {model.total_free_parameters}")
+
+print("\n--- Instantiate concrete model ---")
+
+with timer.section("instance_from_vector"):
+    param_vector = model.physical_values_from_prior_medians
+    instance = model.instance_from_vector(vector=param_vector)
+
+with timer.section("register_pytrees"):
+    _register_model_pytrees(model)
+
+params_tree = jax.tree_util.tree_map(jnp.asarray, instance)
+
+
+# ---------------------------------------------------------------------------
+# Eager baseline — full FitPointDataset (source-plane chi-squared)
+# ---------------------------------------------------------------------------
+
+print("\n--- Eager FitPointDataset (source-plane) ---")
+
+analysis_eager = al.AnalysisPoint(
+    dataset=dataset,
+    solver=solver,
+    fit_positions_cls=al.FitPositionsSource,
+    use_jax=False,
+)
+
+with timer.section("fit_eager"):
+    fit_eager = analysis_eager.fit_from(instance=instance)
+    log_likelihood_ref = float(fit_eager.log_likelihood)
+    figure_of_merit_ref = float(fit_eager.figure_of_merit)
+
+n_eager_repeats = 100
+with timer.section(f"eager_log_likelihood_x{n_eager_repeats}"):
+    for _ in range(n_eager_repeats):
+        analysis_eager.log_likelihood_function(instance=instance)
+eager_per_call = timer.records[-1][1] / n_eager_repeats
+
+print(f"  log_likelihood   = {log_likelihood_ref}")
+print(f"  figure_of_merit  = {figure_of_merit_ref}")
+print(f"  eager per-call   = {eager_per_call:.6f} s")
+
+
+# ===================================================================
+# PART B — Full-pipeline JIT (expected to fail — see module docstring)
+# ===================================================================
+
+print("\n" + "=" * 70)
+print("FULL-PIPELINE JIT (source-plane)")
+print("=" * 70)
+
+analysis_jax = al.AnalysisPoint(
+    dataset=dataset,
+    solver=solver,
+    fit_positions_cls=al.FitPositionsSource,
+    use_jax=True,
+)
+
+
+def full_pipeline_from_params(params_tree):
+    return analysis_jax.log_likelihood_function(instance=params_tree)
+
+
+full_pipeline_jits = False
+full_pipeline_per_call = None
+full_result = None
+full_pipeline_blocker = None
+
+try:
+    _, full_result = jit_profile(
+        full_pipeline_from_params, "full_pipeline", params_tree
+    )
+    full_pipeline_per_call = timer.records[-1][1] / 10
+    full_pipeline_jits = True
+    print(f"  full log_likelihood = {full_result}")
+except jax.errors.TracerArrayConversionError as e:
+    full_pipeline_blocker = (
+        "Grid2DIrregular.grid_2d_via_deflection_grid_from does not propagate xp; "
+        "model_data ends up with _xp=np while holding JAX tracers, so "
+        "squared_distances_to_coordinate_from calls np.square on a tracer."
+    )
+    print(
+        "\n  >>> BLOCKER: full-pipeline source-plane likelihood does NOT JIT.\n"
+        f"  >>> Cause:   {full_pipeline_blocker}\n"
+        "  >>> See module docstring for the proposed library fix."
+    )
+
+
+# ===================================================================
+# PART C — JIT-able prefix: tracer ray-trace of observed positions
+# ===================================================================
+#
+# Even though the full pipeline is blocked, the dominant work in the
+# source-plane likelihood — ray-tracing the observed image positions to
+# the source plane via the tracer's deflection field — IS JIT-traceable
+# when the input/output stay as raw arrays.  We profile that prefix here
+# so the JIT-able portion of the source-plane path is still measured.
+
+print("\n" + "=" * 70)
+print("JIT-ABLE PREFIX: ray-trace observed positions to source plane")
+print("=" * 70)
+
+observed_positions_raw = jnp.array(dataset.positions.array)
+
+
+def ray_trace_to_source_plane(params_tree, positions_raw):
+    """Ray-trace observed image positions to the source plane (raw arrays)."""
+    tracer = al.Tracer(galaxies=list(params_tree.galaxies))
+    grid_in = aa.Grid2DIrregular(values=positions_raw, xp=jnp)
+    deflections = tracer.deflections_yx_2d_from(grid=grid_in, xp=jnp)
+    # Source-plane positions = observed - deflections.
+    return positions_raw - deflections.array
+
+
+_, source_plane_positions = jit_profile(
+    ray_trace_to_source_plane,
+    "raytrace_prefix",
+    params_tree,
+    observed_positions_raw,
+)
+prefix_per_call = timer.records[-1][1] / 10
+
+print(f"  source-plane positions shape: {source_plane_positions.shape}")
+print(f"  source-plane positions value: {np.array(source_plane_positions)}")
+
+
+# ===================================================================
+# PART D — vmap over the JIT-able prefix
+# ===================================================================
+
+print("\n--- vmap over ray-trace prefix ---")
+
+batch_size = 3
+
+batched_params = jax.tree_util.tree_map(
+    lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)),
+    params_tree,
+)
+batched_positions = jnp.broadcast_to(
+    observed_positions_raw, (batch_size, *observed_positions_raw.shape)
+)
+
+vmapped_prefix = jax.jit(jax.vmap(ray_trace_to_source_plane))
+
+with timer.section("vmap_prefix_first_call"):
+    result_vmap = vmapped_prefix(batched_params, batched_positions)
+    block(result_vmap)
+
+n_vmap_repeats = 10
+with timer.section(f"vmap_prefix_steady_x{n_vmap_repeats}"):
+    for _ in range(n_vmap_repeats):
+        result_vmap = vmapped_prefix(batched_params, batched_positions)
+        block(result_vmap)
+
+vmap_batch_time = timer.records[-1][1] / n_vmap_repeats
+vmap_per_call = vmap_batch_time / batch_size
+vmap_speedup = prefix_per_call / vmap_per_call
+
+print(f"  vmap batch={batch_size}: {vmap_batch_time:.6f} s")
+print(f"  vmap per call:         {vmap_per_call:.6f} s")
+print(f"  single JIT per call:   {prefix_per_call:.6f} s")
+print(f"  vmap speedup:          {vmap_speedup:.1f}x faster per ray-trace")
+
+# Eager ray-trace truth — compare against vmap output to lock the prefix.
+eager_grid = aa.Grid2DIrregular(values=np.array(observed_positions_raw))
+eager_deflections = al.Tracer(galaxies=list(instance.galaxies)).deflections_yx_2d_from(
+    grid=eager_grid, xp=np
+)
+eager_source_positions = np.array(observed_positions_raw) - eager_deflections.array
+
+np.testing.assert_allclose(
+    np.array(source_plane_positions),
+    eager_source_positions,
+    rtol=1e-4,
+    err_msg="point_source/source_plane: JIT ray-trace prefix mismatch with eager NumPy",
+)
+np.testing.assert_allclose(
+    np.array(result_vmap),
+    eager_source_positions[None, :, :].repeat(batch_size, axis=0),
+    rtol=1e-4,
+    err_msg="point_source/source_plane: vmap ray-trace prefix mismatch with eager NumPy",
+)
+print("  Eager vs JIT vs vmap (prefix) assertion PASSED")
+
+
+# ===================================================================
+# PART E — Static memory analysis (JIT-able prefix)
+# ===================================================================
+
+print("\n--- Static memory analysis (JIT-able prefix) ---")
+
+lowered_batched = vmapped_prefix.lower(batched_params, batched_positions)
+compiled_batched = lowered_batched.compile()
+mem = compiled_batched.memory_analysis()
+print(f"  Output size:  {mem.output_size_in_bytes / 1024**2:.3f} MB")
+print(f"  Temp size:    {mem.temp_size_in_bytes / 1024**2:.3f} MB")
+print(
+    f"  Total:        "
+    f"{(mem.output_size_in_bytes + mem.temp_size_in_bytes) / 1024**2:.3f} MB"
+)
+
+
+# ===================================================================
+# Summary + outputs
+# ===================================================================
+
+al_version = al.__version__
+
+print("\n" + "=" * 70)
+print(f"JAX LIKELIHOOD SUMMARY — POINT SOURCE SOURCE-PLANE — v{al_version}")
+print("=" * 70)
+print(f"  Dataset:                    {dataset_name}")
+print(f"  Observed image positions:   {n_observed_positions}")
+print(f"  Position noise sigma:       {positions_noise_sigma}")
+print(f"  Free parameters:            {model.total_free_parameters}")
+print(f"  fit_positions_cls:          FitPositionsSource (source-plane chi-squared)")
+print("-" * 70)
+print(
+    f"  Eager full likelihood:      {eager_per_call:.6f} s/call  ({log_likelihood_ref:.6f})"
+)
+if full_pipeline_jits:
+    print(f"  Full pipeline (JIT):        {full_pipeline_per_call:.6f} s/call")
+else:
+    print("  Full pipeline (JIT):        BLOCKED (see module docstring)")
+print(f"  JIT-able prefix (raytrace): {prefix_per_call:.6f} s/call")
+print(f"  vmap prefix per-call (b={batch_size}): {vmap_per_call:.6f} s")
+print(f"  vmap speedup vs single JIT prefix: {vmap_speedup:.1f}x")
+print("=" * 70)
+
+likelihood_summary = {
+    "autolens_version": al_version,
+    "dataset": dataset_name,
+    "fit_positions_cls": "FitPositionsSource",
+    "configuration": {
+        "observed_image_positions": int(n_observed_positions),
+        "positions_noise_sigma": positions_noise_sigma,
+        "free_parameters": int(model.total_free_parameters),
+    },
+    "eager_per_call": eager_per_call,
+    "eager_log_likelihood": log_likelihood_ref,
+    "full_pipeline_jits": full_pipeline_jits,
+    "full_pipeline_blocker": full_pipeline_blocker,
+    "full_pipeline_single_jit": full_pipeline_per_call,
+    "jit_able_prefix": {
+        "name": "ray-trace observed positions to source plane",
+        "per_call": prefix_per_call,
+    },
+    "vmap_prefix": {
+        "batch_size": batch_size,
+        "batch_time": vmap_batch_time,
+        "per_call": vmap_per_call,
+        "speedup_vs_single_jit_prefix": round(vmap_speedup, 1),
+    },
+}
+
+results_dir = _workspace_root / "results" / "likelihood" / "point_source"
+results_dir.mkdir(parents=True, exist_ok=True)
+
+dict_path = results_dir / f"source_plane_summary_v{al_version}.json"
+dict_path.write_text(json.dumps(likelihood_summary, indent=2))
+print(f"\n  Results dict saved to: {dict_path}")
+
+# --- Bar chart ---
+
+labels = [
+    "Eager full likelihood",
+    "JIT-able prefix (raytrace)",
+    f"vmap prefix per-call (batch={batch_size})",
+]
+times = [eager_per_call, prefix_per_call, vmap_per_call]
+colors = ["#8172B3", "#4C72B0", "#55A868"]
+if full_pipeline_jits:
+    labels.insert(1, "Full pipeline (JIT)")
+    times.insert(1, full_pipeline_per_call)
+    colors.insert(1, "#C44E52")
+
+fig, ax = plt.subplots(figsize=(10, 4.5))
+y_pos = range(len(labels))
+bars = ax.barh(y_pos, times, color=colors, edgecolor="white", height=0.6)
+
+for bar, t in zip(bars, times):
+    ax.text(
+        bar.get_width() + max(times) * 0.01,
+        bar.get_y() + bar.get_height() / 2,
+        f"{t:.6f} s",
+        va="center",
+        fontsize=9,
+    )
+
+ax.set_yticks(y_pos)
+ax.set_yticklabels(labels, fontsize=10)
+ax.invert_yaxis()
+ax.set_xlabel("Time per call (s)", fontsize=11)
+fig.suptitle(
+    "Point-Source Likelihood — Source-Plane Chi-Squared",
+    fontsize=12,
+    fontweight="bold",
+)
+title_extra = " | full pipeline JIT BLOCKED" if not full_pipeline_jits else ""
+ax.set_title(
+    f"AutoLens v{al_version}  |  {n_observed_positions} positions  |  "
+    f"{model.total_free_parameters} free params{title_extra}",
+    fontsize=9,
+)
+ax.margins(x=0.20)
+fig.tight_layout()
+
+chart_path = results_dir / f"source_plane_summary_v{al_version}.png"
+fig.savefig(chart_path, dpi=150)
+plt.close(fig)
+print(f"  Bar chart saved to:    {chart_path}")
+
+
+# ===================================================================
+# Regression assertions (eager and full-pipeline JIT)
+# ===================================================================
+#
+# Simulator truth parameters (einstein_radius=1.6, centre=(0.01,0.01),
+# ell_comps=(0.01,0.01), source centre=(0.07,0.07)) + seeded noise
+# (noise_seed=1 in simulators/point_source.py) make the log-likelihood
+# deterministic. Eager numpy and full-pipeline JIT agree to float64.
+EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE = -294.1401881258811
+
+np.testing.assert_allclose(
+    log_likelihood_ref,
+    EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE,
+    rtol=1e-4,
+    err_msg=(
+        f"point_source/source_plane: regression — eager log_likelihood drifted "
+        f"(got {log_likelihood_ref}, expected {EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE})"
+    ),
+)
+print(
+    f"  Eager regression assertion PASSED: log_likelihood matches "
+    f"{EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE:.6f}"
+)
+
+if full_pipeline_jits:
+    np.testing.assert_allclose(
+        float(full_result),
+        EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE,
+        rtol=1e-4,
+        err_msg=(
+            f"point_source/source_plane: regression — JIT log_likelihood drifted "
+            f"(got {float(full_result)}, expected {EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE})"
+        ),
+    )
+    print(
+        f"  JIT regression assertion PASSED: log_likelihood matches "
+        f"{EXPECTED_LOG_LIKELIHOOD_SOURCE_PLANE:.6f}"
+    )
+
+timer.summary()