diff --git a/notebooks/group/simulator.ipynb b/notebooks/group/simulator.ipynb
index 96eff58ed..858cf2998 100644
--- a/notebooks/group/simulator.ipynb
+++ b/notebooks/group/simulator.ipynb
@@ -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."
    ]
   }
  ],
diff --git a/notebooks/guides/tracer.ipynb b/notebooks/guides/tracer.ipynb
index 708d4b65e..ad965b7b4 100644
--- a/notebooks/guides/tracer.ipynb
+++ b/notebooks/guides/tracer.ipynb
@@ -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",
diff --git a/notebooks/guides/using_jax.ipynb b/notebooks/guides/using_jax.ipynb
index e6993b4ef..01841b822 100644
--- a/notebooks/guides/using_jax.ipynb
+++ b/notebooks/guides/using_jax.ipynb
@@ -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",
@@ -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",
diff --git a/notebooks/imaging/simulator.ipynb b/notebooks/imaging/simulator.ipynb
index 05a4adc71..1539a2733 100644
--- a/notebooks/imaging/simulator.ipynb
+++ b/notebooks/imaging/simulator.ipynb
@@ -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
@@ -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",
diff --git a/notebooks/interferometer/simulator.ipynb b/notebooks/interferometer/simulator.ipynb
index 12e18cf07..dcb6c1f4a 100644
--- a/notebooks/interferometer/simulator.ipynb
+++ b/notebooks/interferometer/simulator.ipynb
@@ -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
@@ -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."
    ]
   }
  ],
diff --git a/notebooks/multi_galaxy/simulator.ipynb b/notebooks/multi_galaxy/simulator.ipynb
index 7984952dc..253fceda3 100644
--- a/notebooks/multi_galaxy/simulator.ipynb
+++ b/notebooks/multi_galaxy/simulator.ipynb
@@ -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."
    ]
   }
  ],
diff --git a/scripts/group/simulator.py b/scripts/group/simulator.py
index a4e2146cb..da8f57480 100644
--- a/scripts/group/simulator.py
+++ b/scripts/group/simulator.py
@@ -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.
 """
diff --git a/scripts/guides/tracer.py b/scripts/guides/tracer.py
index dbe534d30..0e1e17411 100644
--- a/scripts/guides/tracer.py
+++ b/scripts/guides/tracer.py
@@ -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`__
 
diff --git a/scripts/guides/using_jax.py b/scripts/guides/using_jax.py
index d5e206101..45a6db84b 100644
--- a/scripts/guides/using_jax.py
+++ b/scripts/guides/using_jax.py
@@ -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.
@@ -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
diff --git a/scripts/imaging/simulator.py b/scripts/imaging/simulator.py
index c5635e2c3..38c1a47a1 100644
--- a/scripts/imaging/simulator.py
+++ b/scripts/imaging/simulator.py
@@ -358,12 +358,9 @@
 
 __JAX Variant (Advanced)__
 
-For an order-of-magnitude speedup on large or repeated simulations (parameter sweeps, mock-data studies, batch
-figure generation), construct the simulator with `use_jax=True` and wrap your call in `@jax.jit`. The simulator
-handles pytree registration internally — you write nothing JAX-specific beyond the decorator.
+For large or repeated simulations (parameter sweeps, mock-data studies, batch figure generation), construct the
+simulator with `use_jax=True` so the image calculation runs through JAX:
 """
-import jax
-
 simulator_jax = al.SimulatorImaging(
     exposure_time=300.0,
     psf=psf,
@@ -372,28 +369,24 @@
     use_jax=True,
 )
 
-
-@jax.jit
-def simulate(tracer):
-    return simulator_jax.via_tracer_from(tracer=tracer, grid=grid)
-
-
 """
-The simulation call below is commented out to avoid adding excessive run time to this script and overwriting the
-dataset output above. Uncomment it to run the JAX simulation, which returns an `Imaging` dataset with `jax.Array`
-data.
-"""
-# dataset_jax = simulate(tracer)
+Call it exactly as above — `simulator_jax.via_tracer_from(tracer=tracer, grid=grid)`. It is not called here to
+avoid overwriting the dataset written above. The returned dataset's `.data.array` is a `numpy.ndarray`.
 
-"""
-The `dataset_jax.data.array` is a `jax.Array`; `aplt.fits_imaging` and the plotters call `numpy.asarray()`
-internally, so saving / plotting works without manual conversion.
+**Wrapping the call in `@jax.jit` does not currently work.** Two separate things stop it:
+
+- **Pytree registration is yours to do, before the first jitted call.** 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 call fails inside autoarray** on array sites that do not yet thread `xp`. Tracked
+  in PyAutoArray; until it is fixed, use the eager call.
 
-Note: eager `simulator_jax.via_tracer_from(tracer, 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.
+`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 why those scripts are fast.
 
-See `scripts/guides/lens_calc.py` for the advanced "JIT-it-yourself" pattern that wraps individual library
-methods like `tracer.image_2d_from` directly.
+See `scripts/guides/using_jax.py` for the full picture, and `scripts/guides/lens_calc.py` for the
+"JIT-it-yourself" pattern that wraps individual library methods like `tracer.image_2d_from` directly.
 
 __Oversampled PSF (Advanced)__
 
diff --git a/scripts/interferometer/simulator.py b/scripts/interferometer/simulator.py
index 3708dcb7c..f334d926a 100644
--- a/scripts/interferometer/simulator.py
+++ b/scripts/interferometer/simulator.py
@@ -335,43 +335,41 @@
 
 __JAX Variant (Advanced)__
 
-For fast repeated interferometer simulations, instantiate the simulator with `use_jax=True` and wrap the call
-in `@jax.jit`. The simulator handles pytree registration internally.
+For fast repeated interferometer simulations, instantiate the simulator with `use_jax=True` so the image
+calculation and transform run through JAX:
 """
-import jax
-import jax.numpy as jnp
-
 simulator_jax = al.SimulatorInterferometer(
     uv_wavelengths=uv_wavelengths,
     exposure_time=300.0,
     noise_sigma=1000.0,
-    transformer_class=al.TransformerNUFFT,  # JAX-native NUFFT (nufftax), works inside @jax.jit
+    transformer_class=al.TransformerNUFFT,  # JAX-native NUFFT (nufftax)
     use_jax=True,
 )
 
+"""
+Call it as `simulator_jax.via_image_from(image=tracer.image_2d_from(grid=grid, xp=jnp))`. It is not called here
+to avoid overwriting the dataset written above.
 
-@jax.jit
-def simulate(tracer):
-    image = tracer.image_2d_from(grid=grid, xp=jnp)
-    return simulator_jax.via_image_from(image=image)
+**Wrapping the call in `@jax.jit` does not currently work.** Two separate things stop it:
 
+- **Pytree registration is yours to do, before the first jitted call** —
+  `autolens.jax.register_tracer_classes(tracer)`. 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.
+- **Even with that, the jitted call fails inside autoarray** on array sites that do not yet thread `xp`. Tracked
+  in PyAutoArray; until it is fixed, use the eager call.
 
-"""
-The simulation call below is commented out to avoid adding excessive run time to this script and overwriting the
-dataset output above. Uncomment it to run the JAX simulation, which returns an `Interferometer` dataset with
-`jax.Array` visibilities.
-"""
-# dataset_jax = simulate(tracer)
+`scripts/point_source/simulator.py` and `scripts/cluster/simulator.py` show the registration step in a
+`PointSolver` context, where `@jax.jit` *does* work.
 
-"""
 Two notes specific to interferometer:
 
 - `TransformerNUFFT` is backed by the JAX-native `nufftax` library (see `__Many Visibilities__` above), so it
   supports `jax.jit` and scales to large UV sets. The legacy pynufft-backed `TransformerNUFFTPyNUFFT` is not
   JAX-traceable; see `autolens_workspace_test/scripts/interferometer/nufft.py` for the parity work. For small
   visibility counts `TransformerDFT` (the simulator default) is also JAX-traceable.
-- Eager `simulator_jax.via_image_from(image)` already runs on JAX without the `@jax.jit` wrap; the JIT only
-  matters for repeated calls.
+- The eager call above works and is the supported route today. Note it returns a dataset whose visibilities are
+  NumPy-backed, not `jax.Array`.
 
-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
+"JIT-it-yourself" pattern applied to individual library methods.
 """
diff --git a/scripts/multi_galaxy/simulator.py b/scripts/multi_galaxy/simulator.py
index 393b7dfdd..187e9a55e 100644
--- a/scripts/multi_galaxy/simulator.py
+++ b/scripts/multi_galaxy/simulator.py
@@ -432,10 +432,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.
 """
diff --git a/workspace_index.json b/workspace_index.json
index 2526a6df8..f1d4343be 100644
--- a/workspace_index.json
+++ b/workspace_index.json
@@ -1376,6 +1376,7 @@
     ],
     "cross_refs": [
       "/guides/advanced/over_sampling.ipynb",
+      "scripts/guides/using_jax.py",
       "scripts/imaging/simulator.py"
     ],
     "notebook": "notebooks/group/simulator.ipynb",
@@ -2443,6 +2444,7 @@
       "scripts/guides/data_structures.py",
       "scripts/guides/galaxies.py",
       "scripts/guides/lens_calc.py",
+      "scripts/guides/using_jax.py",
       "scripts/point_source/simulator.py",
       "start_here.ipynb"
     ],
@@ -2502,14 +2504,12 @@
       "data_structures.py",
       "lens_calc.py",
       "likelihood_function.py",
+      "scripts/cluster/simulator.py",
       "scripts/guides/data_structures.py",
       "scripts/guides/lens_calc.py",
       "scripts/imaging/likelihood_function.py",
-      "scripts/imaging/simulator.py",
       "scripts/interferometer/likelihood_function.py",
-      "scripts/interferometer/simulator.py",
       "scripts/point_source/simulator.py",
-      "simulator.py",
       "start_here.py"
     ],
     "notebook": "notebooks/guides/using_jax.ipynb",
@@ -4286,7 +4286,10 @@
       "imaging/fit.py",
       "imaging/likelihood_function.py",
       "imaging/modeling.py",
-      "scripts/guides/lens_calc.py"
+      "scripts/cluster/simulator.py",
+      "scripts/guides/lens_calc.py",
+      "scripts/guides/using_jax.py",
+      "scripts/point_source/simulator.py"
     ],
     "notebook": "notebooks/imaging/simulator.ipynb",
     "path": "scripts/imaging/simulator.py",
@@ -5368,7 +5371,10 @@
     ],
     "cross_refs": [
       "autolens_workspace_test/scripts/interferometer/nufft.py",
-      "scripts/guides/lens_calc.py"
+      "scripts/cluster/simulator.py",
+      "scripts/guides/lens_calc.py",
+      "scripts/guides/using_jax.py",
+      "scripts/point_source/simulator.py"
     ],
     "notebook": "notebooks/interferometer/simulator.ipynb",
     "path": "scripts/interferometer/simulator.py",
@@ -5964,6 +5970,7 @@
       "multi_galaxy/likelihood_function.py",
       "multi_galaxy/modeling.py",
       "multi_galaxy/start_here.py",
+      "scripts/guides/using_jax.py",
       "scripts/imaging/simulator.py"
     ],
     "notebook": "notebooks/multi_galaxy/simulator.ipynb",