diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index d6694ff5f3..78e20b2c44 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -26,6 +26,7 @@ pip install theano-pymc
This new version of `Theano-PyMC` comes with an experimental JAX backend which, when combined with the new and experimental JAX samplers in PyMC3, can greatly speed up sampling in your model. As this is still very new, please do not use it in production yet but do test it out and let us know if anything breaks and what results you are seeing, especially speed-wise.
### New features
+- New experimental JAX samplers in `pymc3.sample_jax` (see [notebook](https://docs.pymc.io/notebooks/GLM-hierarchical-jax.html) and [#4247](https://github.com/pymc-devs/pymc3/pull/4247)). Requires JAX and either TFP or numpyro.
- Add MLDA, a new stepper for multilevel sampling. MLDA can be used when a hierarchy of approximate posteriors of varying accuracy is available, offering improved sampling efficiency especially in high-dimensional problems and/or where gradients are not available (see [#3926](https://github.com/pymc-devs/pymc3/pull/3926))
- Add Bayesian Additive Regression Trees (BARTs) [#4183](https://github.com/pymc-devs/pymc3/pull/4183))
- Added `pymc3.gp.cov.Circular` kernel for Gaussian Processes on circular domains, e.g. the unit circle (see [#4082](https://github.com/pymc-devs/pymc3/pull/4082)).
diff --git a/docs/source/notebooks/GLM-hierarchical-jax.ipynb b/docs/source/notebooks/GLM-hierarchical-jax.ipynb
new file mode 100644
index 0000000000..c7ccba27f8
--- /dev/null
+++ b/docs/source/notebooks/GLM-hierarchical-jax.ipynb
@@ -0,0 +1,384 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Using JAX for faster sampling\n",
+ "\n",
+ "(c) Thomas Wiecki, 2020\n",
+ "\n",
+ "*Note: These samplers are still experimental.*\n",
+ "\n",
+ "Using the new Theano JAX linker that Brandon Willard has developed, we can compile PyMC3 models to JAX without any change to the PyMC3 code base or any user-level code changes. The way this works is that we take our Theano graph built by PyMC3 and then translate it to JAX primitives. \n",
+ "\n",
+ "Using our Python samplers, this is still a bit slower than the C-code generated by default Theano.\n",
+ "\n",
+ "However, things get really interesting when we also express our samplers in JAX. Here we have used the JAX samplers by NumPyro or TFP. This combining of the samplers was done by [Junpeng Lao](https://twitter.com/junpenglao). \n",
+ "\n",
+ "The reason this is so much faster is that while before in PyMC3, only the logp evaluation was compiled while the samplers where still coded in Python, so for every loop we went back from C to Python. With this approach, the model *and* the sampler are JIT-compiled by JAX and there is no more Python overhead during the whole sampling run. This way we also get sampling on GPUs or TPUs for free.\n",
+ "\n",
+ "This NB requires the master of [Theano-PyMC](https://github.com/pymc-devs/Theano-PyMC), the [pymc3jax branch of PyMC3](https://github.com/pymc-devs/pymc3/tree/pymc3jax), as well as JAX, TFP-nightly and numpyro.\n",
+ "\n",
+ "This is all still highly experimental but extremely promising and just plain amazing.\n",
+ "\n",
+ "As an example we'll use the classic Radon hierarchical model. Note that this model is still very small, I would expect much more massive speed-ups with larger models."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Running on PyMC3 v3.9.3\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/Users/twiecki/projects/pymc/pymc3/sampling_jax.py:20: UserWarning: This module is still highly experimental.\n",
+ " warnings.warn(\"This module is still highly experimental.\")\n"
+ ]
+ }
+ ],
+ "source": [
+ "import arviz as az\n",
+ "import matplotlib.pyplot as plt\n",
+ "import numpy as np\n",
+ "import pandas as pd\n",
+ "import pymc3 as pm\n",
+ "import pymc3.sampling_jax\n",
+ "import theano\n",
+ "\n",
+ "print(f\"Running on PyMC3 v{pm.__version__}\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%config InlineBackend.figure_format = 'retina'\n",
+ "az.style.use(\"arviz-darkgrid\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "data = pd.read_csv(pm.get_data(\"radon.csv\"))\n",
+ "data[\"log_radon\"] = data[\"log_radon\"].astype(theano.config.floatX)\n",
+ "county_names = data.county.unique()\n",
+ "county_idx = data.county_code.values.astype(\"int32\")\n",
+ "\n",
+ "n_counties = len(data.county.unique())"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Unchanged PyMC3 model specification:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "with pm.Model() as hierarchical_model:\n",
+ " # Hyperpriors for group nodes\n",
+ " mu_a = pm.Normal(\"mu_a\", mu=0.0, sigma=100.0)\n",
+ " sigma_a = pm.HalfNormal(\"sigma_a\", 5.0)\n",
+ " mu_b = pm.Normal(\"mu_b\", mu=0.0, sigma=100.0)\n",
+ " sigma_b = pm.HalfNormal(\"sigma_b\", 5.0)\n",
+ "\n",
+ " # Intercept for each county, distributed around group mean mu_a\n",
+ " # Above we just set mu and sd to a fixed value while here we\n",
+ " # plug in a common group distribution for all a and b (which are\n",
+ " # vectors of length n_counties).\n",
+ " a = pm.Normal(\"a\", mu=mu_a, sigma=sigma_a, shape=n_counties)\n",
+ " # Intercept for each county, distributed around group mean mu_a\n",
+ " b = pm.Normal(\"b\", mu=mu_b, sigma=sigma_b, shape=n_counties)\n",
+ "\n",
+ " # Model error\n",
+ " eps = pm.HalfCauchy(\"eps\", 5.0)\n",
+ "\n",
+ " radon_est = a[county_idx] + b[county_idx] * data.floor.values\n",
+ "\n",
+ " # Data likelihood\n",
+ " radon_like = pm.Normal(\"radon_like\", mu=radon_est, sigma=eps, observed=data.log_radon)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Sampling using our old Python NUTS sampler"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "Auto-assigning NUTS sampler...\n",
+ "Initializing NUTS using jitter+adapt_diag...\n",
+ "Multiprocess sampling (2 chains in 2 jobs)\n",
+ "NUTS: [eps, b, a, sigma_b, mu_b, sigma_a, mu_a]\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "