Skip to content

docs: correct false simulator-JAX claims (registration, eager-on-JAX, TransformerNUFFT) #380

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 2 commits into from
Jul 30, 2026
Merged

docs: correct false simulator-JAX claims (registration, eager-on-JAX, TransformerNUFFT) #380

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e7210ca
docs: correct false simulator-JAX claims (registration, eager-on-JAX,…
Jul 30, 2026
fc1fde3
Merge remote-tracking branch 'origin/main' into feature/correct-simul…
Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 7 additions & 6 deletions notebooks/group/simulator.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -600,12 +600,13 @@
"\n",
"__JAX Variant__\n",
"\n",
"Same pattern as `scripts/imaging/simulator.py` `__JAX Variant__`:\n",
"instantiate `al.SimulatorImaging(use_jax=True)` and wrap\n",
"`via_tracer_from` in `@jax.jit`. The simulator handles pytree\n",
"registration internally.\n",
"\n",
"See `scripts/imaging/simulator.py` for the runnable variant block."
"Same as `scripts/imaging/simulator.py` `__JAX Variant (Advanced)__`:\n",
"instantiate `al.SimulatorImaging(use_jax=True)` and call\n",
"`via_tracer_from` eagerly. Wrapping it in `@jax.jit` does not\n",
"currently work \u2014 registration is the caller's job\n",
"(`autolens.jax.register_tracer_classes`) and, even with it, the\n",
"jitted call fails inside autoarray. See that script, or\n",
"`scripts/guides/using_jax.py`, for the detail."
]
}
],
Expand Down
14 changes: 9 additions & 5 deletions notebooks/guides/tracer.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -900,12 +900,16 @@
"\n",
"You access this in two ways.\n",
"\n",
"__1. The implicit path: `Analysis` and `Simulator`__\n",
"__1. The implicit path: `Analysis`__\n",
"\n",
"`AnalysisImaging(use_jax=True)` (the default) and\n",
"`SimulatorImaging(use_jax=True)` both JAX-accelerate the tracer\n",
"internally. Pytree registration runs as a side effect of the first\n",
"`fit_from` / `via_tracer_from` call; you write nothing JAX-specific.\n",
"`AnalysisImaging(use_jax=True)` (the default) JAX-accelerates the tracer\n",
"internally, and pytree registration runs as a side effect of the first\n",
"`fit_from` call \u2014 you write nothing JAX-specific.\n",
"\n",
"`SimulatorImaging(use_jax=True)` also runs the image calculation through\n",
"JAX, but it does *not* register pytrees for you, so it is only implicit\n",
"for the eager call. See `scripts/guides/using_jax.py` before wrapping a\n",
"simulator in `@jax.jit`.\n",
"\n",
"__2. The explicit path: your own `@jax.jit`__\n",
"\n",
Expand Down
36 changes: 20 additions & 16 deletions notebooks/guides/using_jax.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,29 +46,31 @@
"\n",
"Two situations call for it:\n",
"\n",
"1. **Custom simulations.** Pass `use_jax=True` to the simulator constructor and wrap your call in `@jax.jit` when\n",
" you want to render many datasets fast \u2014 parameter sweeps, mock-data studies, batch figure generation.\n",
" For example:\n",
"1. **Custom simulations.** Pass `use_jax=True` to the simulator constructor to run the image calculation through\n",
" JAX, for parameter sweeps, mock-data studies or batch figure generation:\n",
"\n",
" ```python\n",
" import jax\n",
"\n",
" simulator = al.SimulatorImaging(\n",
" exposure_time=300.0, psf=psf, background_sky_level=0.1, use_jax=True\n",
" )\n",
"\n",
" @jax.jit\n",
" def simulate(tracer):\n",
" return simulator.via_tracer_from(tracer=tracer, grid=grid)\n",
" dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n",
" ```\n",
"\n",
" The simulator handles pytree registration internally, so you write nothing JAX-specific beyond the decorator.\n",
" Note that eager `simulator.via_tracer_from(tracer, grid)` (no `@jax.jit`) already runs on JAX and is sufficient\n",
" for one-off simulations \u2014 the `@jax.jit` wrap only pays off when you call the function many times.\n",
" **Wrapping that call in `@jax.jit` does not currently work.** Two things stop it, and it is worth knowing\n",
" which is which:\n",
"\n",
" - **You must register the pytrees yourself first.** Nothing in the library does it for you, and nothing can:\n",
" JAX flattens a jitted function's arguments at trace time, *before* entering the callee, so a simulator that\n",
" registered internally would already be too late. The one-time call is\n",
" `autolens.jax.register_tracer_classes(tracer)`.\n",
" - **Even with that, the jitted simulator call fails inside autoarray** on array sites that do not yet thread\n",
" `xp` \u2014 see PyAutoLabs/PyAutoArray for the tracked issue. Until it is fixed, use the eager call above.\n",
"\n",
" Note the eager call returns a dataset whose `.data.array` is a `numpy.ndarray`, not a `jax.Array`.\n",
"\n",
" The per-dataset-type `simulator.py` scripts (`scripts/imaging/simulator.py`,\n",
" `scripts/interferometer/simulator.py`, `scripts/point_source/simulator.py`) each show the canonical pattern in\n",
" their `__JAX Variant__` section.\n",
" `scripts/point_source/simulator.py` and `scripts/cluster/simulator.py` show the registration step in a\n",
" `PointSolver` context, where `@jax.jit` *does* work and is the reason those scripts are fast.\n",
"\n",
"2. **Custom likelihood functions** that you assemble by hand rather than reaching for `AnalysisImaging`. Same\n",
" shape: `@jax.jit` around your own `def log_likelihood(instance): ...`. The next section works this through.\n",
Expand Down Expand Up @@ -126,8 +128,10 @@
"rather than passed across its boundary. That call is for the `__JIT-ing Library Methods__` case below, where a\n",
"`Tracer` is an argument.\n",
"\n",
"For interferometer data the same shape applies with `al.FitInterferometer`, with one constraint: use\n",
"`TransformerDFT` (the default). `TransformerNUFFT` is not JAX-traceable.\n",
"For interferometer data the same shape applies with `al.FitInterferometer`. Both `TransformerDFT` and the\n",
"nufftax-backed `TransformerNUFFT` are JAX-traceable, so either works; only the legacy pynufft-backed\n",
"`TransformerNUFFTPyNUFFT` is not. Note the defaults differ by class: `Interferometer` (what a fit uses) defaults\n",
"to `TransformerNUFFT`, while `SimulatorInterferometer` defaults to `TransformerDFT`.\n",
"\n",
"**Via `Fitness` \u2014 the production path.** A non-linear search does not call your function; it calls a `Fitness`\n",
"object, which maps a raw parameter vector to a model instance, calls the analysis, and returns the figure of\n",
Expand Down
53 changes: 18 additions & 35 deletions notebooks/imaging/simulator.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -610,47 +610,21 @@
"\n",
"__JAX Variant (Advanced)__\n",
"\n",
"For an order-of-magnitude speedup on large or repeated simulations (parameter sweeps, mock-data studies, batch\n",
"figure generation), construct the simulator with `use_jax=True` and wrap your call in `@jax.jit`. The simulator\n",
"handles pytree registration internally \u2014 you write nothing JAX-specific beyond the decorator."
"For large or repeated simulations (parameter sweeps, mock-data studies, batch figure generation), construct the\n",
"simulator with `use_jax=True` so the image calculation runs through JAX:"
]
},
{
"cell_type": "code",
"metadata": {},
"source": [
"import jax\n",
"\n",
"simulator_jax = al.SimulatorImaging(\n",
" exposure_time=300.0,\n",
" psf=psf,\n",
" background_sky_level=0.1,\n",
" add_poisson_noise_to_data=True,\n",
" use_jax=True,\n",
")\n",
"\n",
"\n",
"@jax.jit\n",
"def simulate(tracer):\n",
" return simulator_jax.via_tracer_from(tracer=tracer, grid=grid)\n"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The simulation call below is commented out to avoid adding excessive run time to this script and overwriting the\n",
"dataset output above. Uncomment it to run the JAX simulation, which returns an `Imaging` dataset with `jax.Array`\n",
"data."
]
},
{
"cell_type": "code",
"metadata": {},
"source": [
"# dataset_jax = simulate(tracer)"
")"
],
"outputs": [],
"execution_count": null
Expand All @@ -659,14 +633,23 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"The `dataset_jax.data.array` is a `jax.Array`; `aplt.fits_imaging` and the plotters call `numpy.asarray()`\n",
"internally, so saving / plotting works without manual conversion.\n",
"Call it exactly as above \u2014 `simulator_jax.via_tracer_from(tracer=tracer, grid=grid)`. It is not called here to\n",
"avoid overwriting the dataset written above. The returned dataset's `.data.array` is a `numpy.ndarray`.\n",
"\n",
"**Wrapping the call in `@jax.jit` does not currently work.** Two separate things stop it:\n",
"\n",
"- **Pytree registration is yours to do, before the first jitted call.** Nothing in the library does it for you,\n",
" and nothing can \u2014 JAX flattens a jitted function's arguments at trace time, before entering the callee, so a\n",
" simulator that registered internally would already be too late. The one-time call is\n",
" `autolens.jax.register_tracer_classes(tracer)`.\n",
"- **Even with that, the jitted call fails inside autoarray** on array sites that do not yet thread `xp`. Tracked\n",
" in PyAutoArray; until it is fixed, use the eager call.\n",
"\n",
"Note: eager `simulator_jax.via_tracer_from(tracer, grid)` (no `@jax.jit`) already runs on JAX and is sufficient\n",
"for one-off simulations. The `@jax.jit` wrap is only beneficial when you call the function many times.\n",
"`scripts/point_source/simulator.py` and `scripts/cluster/simulator.py` show the registration step in a\n",
"`PointSolver` context, where `@jax.jit` *does* work and is why those scripts are fast.\n",
"\n",
"See `scripts/guides/lens_calc.py` for the advanced \"JIT-it-yourself\" pattern that wraps individual library\n",
"methods like `tracer.image_2d_from` directly.\n",
"See `scripts/guides/using_jax.py` for the full picture, and `scripts/guides/lens_calc.py` for the\n",
"\"JIT-it-yourself\" pattern that wraps individual library methods like `tracer.image_2d_from` directly.\n",
"\n",
"__Oversampled PSF (Advanced)__\n",
"\n",
Expand Down
56 changes: 22 additions & 34 deletions notebooks/interferometer/simulator.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -564,48 +564,21 @@
"\n",
"__JAX Variant (Advanced)__\n",
"\n",
"For fast repeated interferometer simulations, instantiate the simulator with `use_jax=True` and wrap the call\n",
"in `@jax.jit`. The simulator handles pytree registration internally."
"For fast repeated interferometer simulations, instantiate the simulator with `use_jax=True` so the image\n",
"calculation and transform run through JAX:"
]
},
{
"cell_type": "code",
"metadata": {},
"source": [
"import jax\n",
"import jax.numpy as jnp\n",
"\n",
"simulator_jax = al.SimulatorInterferometer(\n",
" uv_wavelengths=uv_wavelengths,\n",
" exposure_time=300.0,\n",
" noise_sigma=1000.0,\n",
" transformer_class=al.TransformerNUFFT, # JAX-native NUFFT (nufftax), works inside @jax.jit\n",
" transformer_class=al.TransformerNUFFT, # JAX-native NUFFT (nufftax)\n",
" use_jax=True,\n",
")\n",
"\n",
"\n",
"@jax.jit\n",
"def simulate(tracer):\n",
" image = tracer.image_2d_from(grid=grid, xp=jnp)\n",
" return simulator_jax.via_image_from(image=image)\n"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The simulation call below is commented out to avoid adding excessive run time to this script and overwriting the\n",
"dataset output above. Uncomment it to run the JAX simulation, which returns an `Interferometer` dataset with\n",
"`jax.Array` visibilities."
]
},
{
"cell_type": "code",
"metadata": {},
"source": [
"# dataset_jax = simulate(tracer)"
")"
],
"outputs": [],
"execution_count": null
Expand All @@ -614,16 +587,31 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Call it as `simulator_jax.via_image_from(image=tracer.image_2d_from(grid=grid, xp=jnp))`. It is not called here\n",
"to avoid overwriting the dataset written above.\n",
"\n",
"**Wrapping the call in `@jax.jit` does not currently work.** Two separate things stop it:\n",
"\n",
"- **Pytree registration is yours to do, before the first jitted call** \u2014\n",
" `autolens.jax.register_tracer_classes(tracer)`. Nothing in the library does it for you, and nothing can: JAX\n",
" flattens a jitted function's arguments at trace time, before entering the callee.\n",
"- **Even with that, the jitted call fails inside autoarray** on array sites that do not yet thread `xp`. Tracked\n",
" in PyAutoArray; until it is fixed, use the eager call.\n",
"\n",
"`scripts/point_source/simulator.py` and `scripts/cluster/simulator.py` show the registration step in a\n",
"`PointSolver` context, where `@jax.jit` *does* work.\n",
"\n",
"Two notes specific to interferometer:\n",
"\n",
"- `TransformerNUFFT` is backed by the JAX-native `nufftax` library (see `__Many Visibilities__` above), so it\n",
" supports `jax.jit` and scales to large UV sets. The legacy pynufft-backed `TransformerNUFFTPyNUFFT` is not\n",
" JAX-traceable; see `autolens_workspace_test/scripts/interferometer/nufft.py` for the parity work. For small\n",
" visibility counts `TransformerDFT` (the simulator default) is also JAX-traceable.\n",
"- Eager `simulator_jax.via_image_from(image)` already runs on JAX without the `@jax.jit` wrap; the JIT only\n",
" matters for repeated calls.\n",
"- The eager call above works and is the supported route today. Note it returns a dataset whose visibilities are\n",
" NumPy-backed, not `jax.Array`.\n",
"\n",
"See `scripts/guides/lens_calc.py` for the \"JIT-it-yourself\" pattern applied to individual library methods."
"See `scripts/guides/using_jax.py` for the full picture, and `scripts/guides/lens_calc.py` for the\n",
"\"JIT-it-yourself\" pattern applied to individual library methods."
]
}
],
Expand Down
13 changes: 7 additions & 6 deletions notebooks/multi_galaxy/simulator.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -728,12 +728,13 @@
"\n",
"__JAX Variant__\n",
"\n",
"Same pattern as `scripts/imaging/simulator.py` `__JAX Variant__`:\n",
"instantiate `al.SimulatorImaging(use_jax=True)` and wrap\n",
"`via_tracer_from` in `@jax.jit`. The simulator handles pytree\n",
"registration internally.\n",
"\n",
"See `scripts/imaging/simulator.py` for the runnable variant block."
"Same as `scripts/imaging/simulator.py` `__JAX Variant (Advanced)__`:\n",
"instantiate `al.SimulatorImaging(use_jax=True)` and call\n",
"`via_tracer_from` eagerly. Wrapping it in `@jax.jit` does not\n",
"currently work \u2014 registration is the caller's job\n",
"(`autolens.jax.register_tracer_classes`) and, even with it, the\n",
"jitted call fails inside autoarray. See that script, or\n",
"`scripts/guides/using_jax.py`, for the detail."
]
}
],
Expand Down
13 changes: 7 additions & 6 deletions scripts/group/simulator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -338,10 +338,11 @@

__JAX Variant__

Same pattern as `scripts/imaging/simulator.py` `__JAX Variant__`:
instantiate `al.SimulatorImaging(use_jax=True)` and wrap
`via_tracer_from` in `@jax.jit`. The simulator handles pytree
registration internally.

See `scripts/imaging/simulator.py` for the runnable variant block.
Same as `scripts/imaging/simulator.py` `__JAX Variant (Advanced)__`:
instantiate `al.SimulatorImaging(use_jax=True)` and call
`via_tracer_from` eagerly. Wrapping it in `@jax.jit` does not
currently work — registration is the caller's job
(`autolens.jax.register_tracer_classes`) and, even with it, the
jitted call fails inside autoarray. See that script, or
`scripts/guides/using_jax.py`, for the detail.
"""
14 changes: 9 additions & 5 deletions scripts/guides/tracer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -529,12 +529,16 @@

You access this in two ways.

__1. The implicit path: `Analysis` and `Simulator`__
__1. The implicit path: `Analysis`__

`AnalysisImaging(use_jax=True)` (the default) and
`SimulatorImaging(use_jax=True)` both JAX-accelerate the tracer
internally. Pytree registration runs as a side effect of the first
`fit_from` / `via_tracer_from` call; you write nothing JAX-specific.
`AnalysisImaging(use_jax=True)` (the default) JAX-accelerates the tracer
internally, and pytree registration runs as a side effect of the first
`fit_from` call — you write nothing JAX-specific.

`SimulatorImaging(use_jax=True)` also runs the image calculation through
JAX, but it does *not* register pytrees for you, so it is only implicit
for the eager call. See `scripts/guides/using_jax.py` before wrapping a
simulator in `@jax.jit`.

__2. The explicit path: your own `@jax.jit`__

Expand Down
36 changes: 20 additions & 16 deletions scripts/guides/using_jax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,29 +41,31 @@

Two situations call for it:

1. **Custom simulations.** Pass `use_jax=True` to the simulator constructor and wrap your call in `@jax.jit` when
you want to render many datasets fast — parameter sweeps, mock-data studies, batch figure generation.
For example:
1. **Custom simulations.** Pass `use_jax=True` to the simulator constructor to run the image calculation through
JAX, for parameter sweeps, mock-data studies or batch figure generation:

```python
import jax

simulator = al.SimulatorImaging(
exposure_time=300.0, psf=psf, background_sky_level=0.1, use_jax=True
)

@jax.jit
def simulate(tracer):
return simulator.via_tracer_from(tracer=tracer, grid=grid)
dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)
```

The simulator handles pytree registration internally, so you write nothing JAX-specific beyond the decorator.
Note that eager `simulator.via_tracer_from(tracer, grid)` (no `@jax.jit`) already runs on JAX and is sufficient
for one-off simulations — the `@jax.jit` wrap only pays off when you call the function many times.
**Wrapping that call in `@jax.jit` does not currently work.** Two things stop it, and it is worth knowing
which is which:

- **You must register the pytrees yourself first.** Nothing in the library does it for you, and nothing can:
JAX flattens a jitted function's arguments at trace time, *before* entering the callee, so a simulator that
registered internally would already be too late. The one-time call is
`autolens.jax.register_tracer_classes(tracer)`.
- **Even with that, the jitted simulator call fails inside autoarray** on array sites that do not yet thread
`xp` — see PyAutoLabs/PyAutoArray for the tracked issue. Until it is fixed, use the eager call above.

Note the eager call returns a dataset whose `.data.array` is a `numpy.ndarray`, not a `jax.Array`.

The per-dataset-type `simulator.py` scripts (`scripts/imaging/simulator.py`,
`scripts/interferometer/simulator.py`, `scripts/point_source/simulator.py`) each show the canonical pattern in
their `__JAX Variant__` section.
`scripts/point_source/simulator.py` and `scripts/cluster/simulator.py` show the registration step in a
`PointSolver` context, where `@jax.jit` *does* work and is the reason those scripts are fast.

2. **Custom likelihood functions** that you assemble by hand rather than reaching for `AnalysisImaging`. Same
shape: `@jax.jit` around your own `def log_likelihood(instance): ...`. The next section works this through.
Expand Down Expand Up @@ -121,8 +123,10 @@ def log_likelihood(instance):
rather than passed across its boundary. That call is for the `__JIT-ing Library Methods__` case below, where a
`Tracer` is an argument.

For interferometer data the same shape applies with `al.FitInterferometer`, with one constraint: use
`TransformerDFT` (the default). `TransformerNUFFT` is not JAX-traceable.
For interferometer data the same shape applies with `al.FitInterferometer`. Both `TransformerDFT` and the
nufftax-backed `TransformerNUFFT` are JAX-traceable, so either works; only the legacy pynufft-backed
`TransformerNUFFTPyNUFFT` is not. Note the defaults differ by class: `Interferometer` (what a fit uses) defaults
to `TransformerNUFFT`, while `SimulatorInterferometer` defaults to `TransformerDFT`.

**Via `Fitness` — the production path.** A non-linear search does not call your function; it calls a `Fitness`
object, which maps a raw parameter vector to a model instance, calls the analysis, and returns the figure of
Expand Down
Loading
Loading