Add JAX with CUDA recipe (#2474)

* Add JAX with CUDA recipe

Closes #2471

* Re-trigger CI

* Use canonical JAX docs URL to avoid linkcheck redirect warning

---------

Co-authored-by: Ayaz Salikhov <mathbunnyru@users.noreply.github.com>

Author: goelakash

docs/using/recipe_code/jax_cuda.dockerfile

@@ -0,0 +1,6 @@
+ARG BASE_IMAGE=quay.io/jupyter/base-notebook
+FROM $BASE_IMAGE
+
+RUN pip install -v --no-cache-dir 'jax[cuda13]' && \
+    fix-permissions "${CONDA_DIR}" && \
+    fix-permissions "/home/${NB_USER}"

docs/using/recipes.md

@@ -121,6 +121,16 @@ You can find the original Jupyter Notebook extension [here](https://github.com/d
 :language: docker
 ```
 
+## JAX with CUDA
+
+[JAX](https://docs.jax.dev/en/latest/) GPU support via the bundled NVIDIA CUDA wheels.
+Run with `--gpus all` and a host NVIDIA driver new enough for CUDA 13.
+For CUDA 12 instead, swap `jax[cuda13]` for `jax[cuda12]`.
+
+```{literalinclude} recipe_code/jax_cuda.dockerfile
+:language: docker
+```
+
 ## Running behind an nginx proxy
 
 ```{warning}