From c1fca4801eaed1d34d3c59686042dad03d0fd6ac Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 25 May 2026 10:55:41 +0100 Subject: [PATCH] docs: remove Array2D.native jit-blocked caveats from JAX variant blocks The array_2d_via_indexes_from jit-safety fix (PyAutoArray#339) makes Array2D.native jit-traceable. Remove the "currently blocked" notes from simulator JAX variant blocks. Co-Authored-By: Claude Opus 4.7 (1M context) --- scripts/imaging/simulator.py | 12 +++--------- scripts/interferometer/simulator.py | 4 +--- 2 files changed, 4 insertions(+), 12 deletions(-) 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.