diff --git a/scripts/imaging/simulator.py b/scripts/imaging/simulator.py index ca644334..45081114 100644 --- a/scripts/imaging/simulator.py +++ b/scripts/imaging/simulator.py @@ -216,15 +216,9 @@ def simulate(galaxies): plotters call `numpy.asarray()` internally, so saving / plotting works without manual conversion. -Two notes: - -- Eager `simulator_jax.via_galaxies_from(galaxies, grid)` (no `@jax.jit`) - already runs on JAX and is sufficient for one-off simulations. The - `@jax.jit` wrap is only beneficial when you call the function many times. -- The `@jax.jit` wrap is currently blocked by a pre-existing `Array2D.native` - jit-incompatibility in autoarray's slim/native reshape. Eager JAX works - today; the `@jax.jit` shown above will work once a separate refactor - task lands. +Note: eager `simulator_jax.via_galaxies_from(galaxies, grid)` (no `@jax.jit`) +already runs on JAX and is sufficient for one-off simulations. The +`@jax.jit` wrap is only beneficial when you call the function many times. See `scripts/guides/api/data_structures.py` for the broader "JIT-it- yourself" pattern. diff --git a/scripts/interferometer/simulator.py b/scripts/interferometer/simulator.py index 6115028c..99bfa303 100644 --- a/scripts/interferometer/simulator.py +++ b/scripts/interferometer/simulator.py @@ -222,9 +222,7 @@ def simulate(galaxies): `nufftax` replacement is a research path (see `autolens_workspace_test/scripts/interferometer/nufft.py`). - Eager `simulator_jax.via_image_from(image)` already runs on JAX without - the `@jax.jit` wrap; the JIT wrap is currently blocked by a pre-existing - `Array2D.native` autoarray limitation. Eager use works today; the - `@jax.jit` wrap shown above will work once a separate refactor lands. + the `@jax.jit` wrap; the JIT only matters for repeated calls. See `scripts/guides/api/data_structures.py` for the broader "JIT-it- yourself" pattern.