From baa18b2b4a30e0b418481b9cd3c8cbdfa53b18a0 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 19 Aug 2026 14:42:27 -0400 Subject: [PATCH] =?UTF-8?q?docs:=20JAX=20installs=20by=20default=20?= =?UTF-8?q?=E2=80=94=20retire=20[jax]-extra=20install=20prose=20(PyAutoLen?= =?UTF-8?q?s#702)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 --- markdown/group/start_here.md | 2 +- markdown/imaging/start_here.md | 4 ++-- markdown/interferometer/start_here.md | 4 ++-- markdown/multi_dataset/start_here.md | 3 +-- markdown/point_source/start_here.md | 6 +++--- markdown/start_here.md | 7 +++---- notebooks/group/start_here.ipynb | 2 +- notebooks/guides/using_jax.ipynb | 4 ++-- notebooks/imaging/start_here.ipynb | 4 ++-- notebooks/interferometer/start_here.ipynb | 4 ++-- notebooks/multi_dataset/start_here.ipynb | 3 +-- notebooks/multi_galaxy/start_here.ipynb | 4 ++-- notebooks/point_source/start_here.ipynb | 6 +++--- requirements.txt | 2 +- scripts/group/start_here.py | 2 +- scripts/guides/using_jax.py | 4 ++-- scripts/imaging/start_here.py | 4 ++-- scripts/interferometer/start_here.py | 4 ++-- scripts/multi_dataset/start_here.py | 3 +-- scripts/multi_galaxy/start_here.py | 4 ++-- scripts/point_source/start_here.py | 6 +++--- start_here.ipynb | 7 +++---- start_here.py | 7 +++---- 23 files changed, 45 insertions(+), 51 deletions(-) diff --git a/markdown/group/start_here.md b/markdown/group/start_here.md index cbeffc745..f2fe85af8 100644 --- a/markdown/group/start_here.md +++ b/markdown/group/start_here.md @@ -40,7 +40,7 @@ __Contents__ __JAX__ PyAutoLens runs group-scale model-fits on JAX by default — `al.AnalysisImaging` -auto-enables `use_jax=True` if you installed `autolens[jax]`. Group fits +auto-enables `use_jax=True` (JAX installs with `autolens` itself). Group fits benefit substantially from GPU acceleration (the multi-galaxy deflection sum is the dominant cost). Expect 5-30 minutes on GPU vs hours on pure NumPy. diff --git a/markdown/imaging/start_here.md b/markdown/imaging/start_here.md index 01283e77d..747aeed34 100644 --- a/markdown/imaging/start_here.md +++ b/markdown/imaging/start_here.md @@ -36,8 +36,8 @@ __Contents__ __JAX__ -PyAutoLens runs imaging model-fits on JAX by default. If you installed -`autolens[jax]`, the `al.AnalysisImaging(dataset=dataset)` line below +PyAutoLens runs imaging model-fits on JAX by default (JAX installs with +`autolens` itself) — the `al.AnalysisImaging(dataset=dataset)` line below auto-enables `use_jax=True`; expect 10-30 minutes on CPU, 1-10 minutes on GPU, vs 1-2 hours on pure NumPy for a typical lens. If you do not have a GPU locally, Google Colab provides free GPUs. diff --git a/markdown/interferometer/start_here.md b/markdown/interferometer/start_here.md index d84100548..c876bb9c0 100644 --- a/markdown/interferometer/start_here.md +++ b/markdown/interferometer/start_here.md @@ -34,8 +34,8 @@ __Contents__ __JAX__ -PyAutoLens runs interferometer model-fits on JAX by default. If you -installed `autolens[jax]`, `al.AnalysisInterferometer(dataset=dataset)` +PyAutoLens runs interferometer model-fits on JAX by default (JAX installs +with `autolens` itself) — `al.AnalysisInterferometer(dataset=dataset)` below auto-enables `use_jax=True`. Use `TransformerDFT` (the default in this script) under JAX — `TransformerNUFFT` (pynufft) is faster on large UV sets but is not JAX-traceable; the `nufftax` replacement (see the diff --git a/markdown/multi_dataset/start_here.md b/markdown/multi_dataset/start_here.md index c127d8f40..ab28e17ba 100644 --- a/markdown/multi_dataset/start_here.md +++ b/markdown/multi_dataset/start_here.md @@ -45,8 +45,7 @@ PyAutoLens runs multi-wavelength model-fits on JAX by default. The per-band `al.AnalysisImaging(dataset=dataset, use_jax=True)` instances below auto-enable JAX, and the `af.FactorGraphModel(*analysis_factor_list, use_jax=True)` further down stitches them together with JAX-aware -broadcasting. If you installed `autolens[jax]`, expect 1-10 minutes per -band on GPU vs hours on pure NumPy. +broadcasting. Expect 1-10 minutes per band on GPU vs hours on pure NumPy. For the broader JAX principles, see the top-level `autolens_workspace/start_here.py` `__JAX__` section. Per-band diff --git a/markdown/point_source/start_here.md b/markdown/point_source/start_here.md index fc0bcc1ae..0ec5c6591 100644 --- a/markdown/point_source/start_here.md +++ b/markdown/point_source/start_here.md @@ -44,9 +44,9 @@ __Contents__ __JAX__ -PyAutoLens runs point-source model-fits on JAX by default. `AnalysisPoint` -auto-enables `use_jax=True` if you installed `autolens[jax]`; the search -driver wraps the likelihood in `jax.vmap(jax.jit(...))`. +PyAutoLens runs point-source model-fits on JAX by default (JAX installs +with `autolens` itself). `AnalysisPoint` auto-enables `use_jax=True`; the +search driver wraps the likelihood in `jax.vmap(jax.jit(...))`. For the broader JAX principles see `autolens_workspace/start_here.py` `__JAX__`. For the most user-impactful piece — the `PointSolver(use_jax=True)` diff --git a/markdown/start_here.md b/markdown/start_here.md index 1c294923c..89b7e32a6 100644 --- a/markdown/start_here.md +++ b/markdown/start_here.md @@ -320,10 +320,9 @@ __JAX__ array library with GPU support and just-in-time compilation). JAX makes lens modeling 10-100x faster on large grids — sometimes more on GPU. -You do not have to do anything to use it. If you installed `autolens` with -the JAX extra (`pip install autolens[jax]` on Python 3.11+), the analysis -objects you'll meet in the `__Lens Modeling__` section below use JAX -automatically. The first time a model-fit starts you'll see a one-time log +You do not have to do anything to use it. JAX installs with `autolens` by +default (`pip install autolens`), and the analysis objects you'll meet in +the `__Lens Modeling__` section below use JAX automatically. The first time a model-fit starts you'll see a one-time log line like `JAX: Applying vmap and jit to likelihood function -- may take a few seconds.` — that's JAX compiling the likelihood function, after which every evaluation re-uses the compiled code. If JAX is not installed, the diff --git a/notebooks/group/start_here.ipynb b/notebooks/group/start_here.ipynb index e7e062cf4..2bc518444 100644 --- a/notebooks/group/start_here.ipynb +++ b/notebooks/group/start_here.ipynb @@ -54,7 +54,7 @@ "__JAX__\n", "\n", "PyAutoLens runs group-scale model-fits on JAX by default \u2014 `al.AnalysisImaging`\n", - "auto-enables `use_jax=True` if you installed `autolens[jax]`. Group fits\n", + "auto-enables `use_jax=True` (JAX installs with `autolens` itself). Group fits\n", "benefit substantially from GPU acceleration (the multi-galaxy deflection\n", "sum is the dominant cost). Expect 5-30 minutes on GPU vs hours on pure\n", "NumPy.\n", diff --git a/notebooks/guides/using_jax.ipynb b/notebooks/guides/using_jax.ipynb index 695b312c0..71db34089 100644 --- a/notebooks/guides/using_jax.ipynb +++ b/notebooks/guides/using_jax.ipynb @@ -11,8 +11,8 @@ "just-in-time compilation). JAX makes lens modeling 10-100x faster on large grids \u2014 sometimes more on GPU \u2014 so the\n", "library is built to use it automatically wherever it helps.\n", "\n", - "The `start_here.py` introduction covers the one thing every user needs to know: install the JAX extra\n", - "(`pip install autolens[jax]` on Python 3.11+) and lens modeling uses JAX automatically. This guide covers the\n", + "The `start_here.py` introduction covers the one thing every user needs to know: JAX installs with\n", + "`autolens` by default (`pip install autolens`) and lens modeling uses it automatically. This guide covers the\n", "technical detail behind that, and the situations where you interact with JAX directly.\n", "\n", "Unlike most guides, the JAX recipes below are **executable code cells this script runs** \u2014 they are verified\n", diff --git a/notebooks/imaging/start_here.ipynb b/notebooks/imaging/start_here.ipynb index bb2feabee..20c3837c4 100644 --- a/notebooks/imaging/start_here.ipynb +++ b/notebooks/imaging/start_here.ipynb @@ -38,8 +38,8 @@ "\n", "__JAX__\n", "\n", - "PyAutoLens runs imaging model-fits on JAX by default. If you installed\n", - "`autolens[jax]`, the `al.AnalysisImaging(dataset=dataset)` line below\n", + "PyAutoLens runs imaging model-fits on JAX by default (JAX installs with\n", + "`autolens` itself) \u2014 the `al.AnalysisImaging(dataset=dataset)` line below\n", "auto-enables `use_jax=True`. JAX is what makes the gradient-based search\n", "used below possible at all \u2014 it supplies the derivatives of the likelihood\n", "and evaluates all of the search's parallel starts in one batched call. If\n", diff --git a/notebooks/interferometer/start_here.ipynb b/notebooks/interferometer/start_here.ipynb index bcb867716..e22a89359 100644 --- a/notebooks/interferometer/start_here.ipynb +++ b/notebooks/interferometer/start_here.ipynb @@ -36,8 +36,8 @@ "\n", "__JAX__\n", "\n", - "PyAutoLens runs interferometer model-fits on JAX by default. If you\n", - "installed `autolens[jax]`, `al.AnalysisInterferometer(dataset=dataset)`\n", + "PyAutoLens runs interferometer model-fits on JAX by default (JAX installs\n", + "with `autolens` itself) \u2014 `al.AnalysisInterferometer(dataset=dataset)`\n", "below auto-enables `use_jax=True`. Use `TransformerDFT` (the default in\n", "this script) under JAX \u2014 `TransformerNUFFT` (pynufft) is faster on large\n", "UV sets but is not JAX-traceable; the `nufftax` replacement (see the\n", diff --git a/notebooks/multi_dataset/start_here.ipynb b/notebooks/multi_dataset/start_here.ipynb index dfa20859b..c3cebc67b 100644 --- a/notebooks/multi_dataset/start_here.ipynb +++ b/notebooks/multi_dataset/start_here.ipynb @@ -47,8 +47,7 @@ "per-band `al.AnalysisImaging(dataset=dataset, use_jax=True)` instances\n", "below auto-enable JAX, and the `af.FactorGraphModel(*analysis_factor_list,\n", "use_jax=True)` further down stitches them together with JAX-aware\n", - "broadcasting. If you installed `autolens[jax]`, expect 1-10 minutes per\n", - "band on GPU vs hours on pure NumPy.\n", + "broadcasting. Expect 1-10 minutes per band on GPU vs hours on pure NumPy.\n", "\n", "For the broader JAX principles, see the top-level\n", "`autolens_workspace/start_here.py` `__JAX__` section. Per-band\n", diff --git a/notebooks/multi_galaxy/start_here.ipynb b/notebooks/multi_galaxy/start_here.ipynb index a783d5f29..85a925459 100644 --- a/notebooks/multi_galaxy/start_here.ipynb +++ b/notebooks/multi_galaxy/start_here.ipynb @@ -71,8 +71,8 @@ "\n", "__JAX__\n", "\n", - "PyAutoLens runs multi-galaxy model-fits on JAX by default \u2014 `al.AnalysisImaging` auto-enables `use_jax=True` if\n", - "you installed `autolens[jax]`. The multi-galaxy deflection sum vectorises cleanly, so fits benefit substantially\n", + "PyAutoLens runs multi-galaxy model-fits on JAX by default \u2014 `al.AnalysisImaging` auto-enables `use_jax=True`\n", + "(JAX installs with `autolens` itself). The multi-galaxy deflection sum vectorises cleanly, so fits benefit substantially\n", "from GPU acceleration. Expect ~10-20 minutes on a GPU for this example.\n", "\n", "For the broader JAX principles, see the top-level `autolens_workspace/start_here.py` `__JAX__` section.\n", diff --git a/notebooks/point_source/start_here.ipynb b/notebooks/point_source/start_here.ipynb index c3e8a661a..e75c8b757 100644 --- a/notebooks/point_source/start_here.ipynb +++ b/notebooks/point_source/start_here.ipynb @@ -51,9 +51,9 @@ "\n", "__JAX__\n", "\n", - "PyAutoLens runs point-source model-fits on JAX by default. `AnalysisPoint`\n", - "auto-enables `use_jax=True` if you installed `autolens[jax]`; the search\n", - "driver wraps the likelihood in `jax.vmap(jax.jit(...))`.\n", + "PyAutoLens runs point-source model-fits on JAX by default (JAX installs\n", + "with `autolens` itself). `AnalysisPoint` auto-enables `use_jax=True`; the\n", + "search driver wraps the likelihood in `jax.vmap(jax.jit(...))`.\n", "\n", "For the broader JAX principles see `autolens_workspace/start_here.py`\n", "`__JAX__`. For the most user-impactful piece \u2014 the `PointSolver(use_jax=True)`\n", diff --git a/requirements.txt b/requirements.txt index 1fef8db74..4f1ea1dda 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,2 @@ -autolens[jax] +autolens numba diff --git a/scripts/group/start_here.py b/scripts/group/start_here.py index b884daefd..1f61d29a5 100644 --- a/scripts/group/start_here.py +++ b/scripts/group/start_here.py @@ -49,7 +49,7 @@ __JAX__ PyAutoLens runs group-scale model-fits on JAX by default — `al.AnalysisImaging` -auto-enables `use_jax=True` if you installed `autolens[jax]`. Group fits +auto-enables `use_jax=True` (JAX installs with `autolens` itself). Group fits benefit substantially from GPU acceleration (the multi-galaxy deflection sum is the dominant cost). Expect 5-30 minutes on GPU vs hours on pure NumPy. diff --git a/scripts/guides/using_jax.py b/scripts/guides/using_jax.py index 9dc28494f..0e9b58f02 100644 --- a/scripts/guides/using_jax.py +++ b/scripts/guides/using_jax.py @@ -6,8 +6,8 @@ just-in-time compilation). JAX makes lens modeling 10-100x faster on large grids — sometimes more on GPU — so the library is built to use it automatically wherever it helps. -The `start_here.py` introduction covers the one thing every user needs to know: install the JAX extra -(`pip install autolens[jax]` on Python 3.11+) and lens modeling uses JAX automatically. This guide covers the +The `start_here.py` introduction covers the one thing every user needs to know: JAX installs with +`autolens` by default (`pip install autolens`) and lens modeling uses it automatically. This guide covers the technical detail behind that, and the situations where you interact with JAX directly. Unlike most guides, the JAX recipes below are **executable code cells this script runs** — they are verified diff --git a/scripts/imaging/start_here.py b/scripts/imaging/start_here.py index 82cddffb2..df7262d8f 100644 --- a/scripts/imaging/start_here.py +++ b/scripts/imaging/start_here.py @@ -33,8 +33,8 @@ __JAX__ -PyAutoLens runs imaging model-fits on JAX by default. If you installed -`autolens[jax]`, the `al.AnalysisImaging(dataset=dataset)` line below +PyAutoLens runs imaging model-fits on JAX by default (JAX installs with +`autolens` itself) — the `al.AnalysisImaging(dataset=dataset)` line below auto-enables `use_jax=True`. JAX is what makes the gradient-based search used below possible at all — it supplies the derivatives of the likelihood and evaluates all of the search's parallel starts in one batched call. If diff --git a/scripts/interferometer/start_here.py b/scripts/interferometer/start_here.py index bb41da5d2..aa70c0f11 100644 --- a/scripts/interferometer/start_here.py +++ b/scripts/interferometer/start_here.py @@ -31,8 +31,8 @@ __JAX__ -PyAutoLens runs interferometer model-fits on JAX by default. If you -installed `autolens[jax]`, `al.AnalysisInterferometer(dataset=dataset)` +PyAutoLens runs interferometer model-fits on JAX by default (JAX installs +with `autolens` itself) — `al.AnalysisInterferometer(dataset=dataset)` below auto-enables `use_jax=True`. Use `TransformerDFT` (the default in this script) under JAX — `TransformerNUFFT` (pynufft) is faster on large UV sets but is not JAX-traceable; the `nufftax` replacement (see the diff --git a/scripts/multi_dataset/start_here.py b/scripts/multi_dataset/start_here.py index 63aa92c50..703546e89 100644 --- a/scripts/multi_dataset/start_here.py +++ b/scripts/multi_dataset/start_here.py @@ -42,8 +42,7 @@ per-band `al.AnalysisImaging(dataset=dataset, use_jax=True)` instances below auto-enable JAX, and the `af.FactorGraphModel(*analysis_factor_list, use_jax=True)` further down stitches them together with JAX-aware -broadcasting. If you installed `autolens[jax]`, expect 1-10 minutes per -band on GPU vs hours on pure NumPy. +broadcasting. Expect 1-10 minutes per band on GPU vs hours on pure NumPy. For the broader JAX principles, see the top-level `autolens_workspace/start_here.py` `__JAX__` section. Per-band diff --git a/scripts/multi_galaxy/start_here.py b/scripts/multi_galaxy/start_here.py index d06083273..da5cbbeb0 100644 --- a/scripts/multi_galaxy/start_here.py +++ b/scripts/multi_galaxy/start_here.py @@ -66,8 +66,8 @@ __JAX__ -PyAutoLens runs multi-galaxy model-fits on JAX by default — `al.AnalysisImaging` auto-enables `use_jax=True` if -you installed `autolens[jax]`. The multi-galaxy deflection sum vectorises cleanly, so fits benefit substantially +PyAutoLens runs multi-galaxy model-fits on JAX by default — `al.AnalysisImaging` auto-enables `use_jax=True` +(JAX installs with `autolens` itself). The multi-galaxy deflection sum vectorises cleanly, so fits benefit substantially from GPU acceleration. Expect ~10-20 minutes on a GPU for this example. For the broader JAX principles, see the top-level `autolens_workspace/start_here.py` `__JAX__` section. diff --git a/scripts/point_source/start_here.py b/scripts/point_source/start_here.py index 11a522509..f288e7830 100644 --- a/scripts/point_source/start_here.py +++ b/scripts/point_source/start_here.py @@ -46,9 +46,9 @@ __JAX__ -PyAutoLens runs point-source model-fits on JAX by default. `AnalysisPoint` -auto-enables `use_jax=True` if you installed `autolens[jax]`; the search -driver wraps the likelihood in `jax.vmap(jax.jit(...))`. +PyAutoLens runs point-source model-fits on JAX by default (JAX installs +with `autolens` itself). `AnalysisPoint` auto-enables `use_jax=True`; the +search driver wraps the likelihood in `jax.vmap(jax.jit(...))`. For the broader JAX principles see `autolens_workspace/start_here.py` `__JAX__`. For the most user-impactful piece — the `PointSolver(use_jax=True)` diff --git a/start_here.ipynb b/start_here.ipynb index 6a8e2e73e..27582deff 100644 --- a/start_here.ipynb +++ b/start_here.ipynb @@ -342,10 +342,9 @@ "array library with GPU support and just-in-time compilation). JAX makes\n", "lens modeling 10-100x faster on large grids \u2014 sometimes more on GPU.\n", "\n", - "You do not have to do anything to use it. If you installed `autolens` with\n", - "the JAX extra (`pip install autolens[jax]` on Python 3.11+), the analysis\n", - "objects you'll meet in the `__Lens Modeling__` section below use JAX\n", - "automatically. The first time a model-fit starts you'll see a one-time log\n", + "You do not have to do anything to use it. JAX installs with `autolens` by\n", + "default (`pip install autolens`), and the analysis objects you'll meet in\n", + "the `__Lens Modeling__` section below use JAX automatically. The first time a model-fit starts you'll see a one-time log\n", "line like `JAX: Applying vmap and jit to likelihood function -- may take a\n", "few seconds.` \u2014 that's JAX compiling the likelihood function, after which\n", "every evaluation re-uses the compiled code. If JAX is not installed, the\n", diff --git a/start_here.py b/start_here.py index 7a98fbef8..02f6968ef 100644 --- a/start_here.py +++ b/start_here.py @@ -216,10 +216,9 @@ array library with GPU support and just-in-time compilation). JAX makes lens modeling 10-100x faster on large grids — sometimes more on GPU. -You do not have to do anything to use it. If you installed `autolens` with -the JAX extra (`pip install autolens[jax]` on Python 3.11+), the analysis -objects you'll meet in the `__Lens Modeling__` section below use JAX -automatically. The first time a model-fit starts you'll see a one-time log +You do not have to do anything to use it. JAX installs with `autolens` by +default (`pip install autolens`), and the analysis objects you'll meet in +the `__Lens Modeling__` section below use JAX automatically. The first time a model-fit starts you'll see a one-time log line like `JAX: Applying vmap and jit to likelihood function -- may take a few seconds.` — that's JAX compiling the likelihood function, after which every evaluation re-uses the compiled code. If JAX is not installed, the