docs: JAX installs by default — retire [jax]-extra install prose (PyAutoLens#702)

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.

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.

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

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

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)`

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

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",

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",

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",

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",

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",

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",

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",

requirements.txt

@@ -1,2 +1,2 @@
-autolens[jax]
+autolens
 numba

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.

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

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

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

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

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.

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)`

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",

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