Conversation
…(task 7/9)
Create autogalaxy interferometer jax_grad scripts under a new
jax_grad/interferometer/ subfolder. Autolens has no interferometer
jax_grad reference today, so these are greenfield ports modelled on
the matching jax_likelihood_functions/interferometer/{lp,mge}.py
scripts.
Each script wraps the AnalysisInterferometer log-likelihood in
jax.value_and_grad and asserts the gradient is finite, has the
expected free-parameter shape, and is not all-zero.
The jax_grad/ env_vars override added in task 6 (PR #29) already
covers jax_grad/interferometer/ — no env_vars.yaml change needed.
Closes #30.
This was referenced May 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Create
scripts/jax_grad/interferometer/{lp.py, mge.py}from scratch — autolens has no interferometerjax_gradreference today. Each script runsjax.value_and_gradon the autogalaxyAnalysisInterferometerlikelihood path and asserts the gradient is finite, the shape matchesmodel.total_free_parameters, and the gradient is not all-zero.This is task 7/9 of the autogalaxy_workspace_test parity epic (#5).
Scripts Changed
scripts/jax_grad/interferometer/__init__.py— new (empty namespace marker).scripts/jax_grad/interferometer/lp.py— new. Single galaxy with a Sersic bulge ondataset/interferometer/jax_testwith a(256,256)pixel_scales=0.1radius=3.0real-space mask andTransformerDFT(mirrorsjax_likelihood_functions/interferometer/lp.py). Body usesjax.value_and_grad(fitness.call)and the four canonical assertions. Local run: 6.9s, gradient shape(7,).scripts/jax_grad/interferometer/mge.py— new. Same dataset/mask/transformer; model swapped to a single-basis MGE (mge_model_from, 20 gaussians,centre_prior_is_uniform=True) on a single galaxy. Same body and assertions. Local run: 12.5s, gradient shape(4,).smoke_tests.txt— appended both new scripts immediately afterjax_grad/imaging/mge.py.Notes
Layout divergence from autolens (continued). Autolens has flat top-level
jax_grad/imaging_*.pyand no interferometerjax_gradscripts at all. This PR continues the subfolder convention established in PR #29 (task 6) —jax_grad/{imaging,interferometer,multi}/on autogalaxy. The retrofit question for autolens was raised in PR #29; suggest filing the follow-up after task 8 lands so the migration covers all three subdirectories at once.lp.Sersicvslp_linear.Sersic. Task 6 (jax_grad/imaging/lp.py) usedlp_linear.Sersic. This task uses plainlp.Sersicto match the validatedjax_likelihood_functions/interferometer/lp.pypattern. The interferometer + JAX inversion path differs from imaging; sticking with the proven setup.No env_vars.yaml change. The
jax_grad/substring override added in PR #29 (unsetPYAUTO_SMALL_DATASETS+PYAUTO_DISABLE_JAX) already matchesjax_grad/interferometer/.Test Plan
.github/scripts/run_smoke.pywith the env override applied:jax_grad/interferometer/lp.py— exit 0 in 6.9s, grad shape(7,)jax_grad/interferometer/mge.py— exit 0 in 12.5s, grad shape(4,)🤖 Generated with Claude Code
Closes #30.