diff --git a/pyproject.toml b/pyproject.toml
index 06168d9b..4f9f8503 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -101,6 +101,7 @@ notebooks = [
     "black",
     "ipympl",
     "pandas",
+    "tensorwaves[jax,pwa]",
 ]
 numba = [
     "numba",
diff --git a/uv.lock b/uv.lock
index 393aca0b..dce68691 100644
--- a/uv.lock
+++ b/uv.lock
@@ -4193,12 +4193,20 @@ dev = [
     { name = "types-tqdm" },
 ]
 doc = [
+    { name = "ampform" },
     { name = "black" },
+    { name = "graphviz" },
     { name = "ipympl" },
+    { name = "jax" },
+    { name = "jaxlib" },
     { name = "jupyter" },
     { name = "matplotlib" },
     { name = "myst-nb" },
+    { name = "numba" },
+    { name = "numpy" },
     { name = "pandas" },
+    { name = "phasespace", extra = ["tf"] },
+    { name = "scipy" },
     { name = "sphinx" },
     { name = "sphinx-api-relink" },
     { name = "sphinx-book-theme" },
@@ -4209,6 +4217,7 @@ doc = [
     { name = "sphinx-thebe" },
     { name = "sphinx-togglebutton" },
     { name = "sphobjinv" },
+    { name = "tensorflow" },
 ]
 jax = [
     { name = "ampform" },
@@ -4222,8 +4231,12 @@ jax = [
     { name = "tensorflow" },
 ]
 jupyter = [
+    { name = "ampform" },
     { name = "black" },
+    { name = "graphviz" },
     { name = "ipympl" },
+    { name = "jax" },
+    { name = "jaxlib" },
     { name = "jupyter" },
     { name = "jupyterlab" },
     { name = "jupyterlab-code-formatter" },
@@ -4232,9 +4245,13 @@ jupyter = [
     { name = "jupyterlab-myst" },
     { name = "matplotlib" },
     { name = "myst-nb" },
+    { name = "numba" },
+    { name = "numpy" },
     { name = "pandas" },
+    { name = "phasespace", extra = ["tf"] },
     { name = "python-lsp-ruff" },
     { name = "python-lsp-server", extra = ["rope"] },
+    { name = "scipy" },
     { name = "sphinx" },
     { name = "sphinx-api-relink" },
     { name = "sphinx-book-theme" },
@@ -4245,11 +4262,21 @@ jupyter = [
     { name = "sphinx-thebe" },
     { name = "sphinx-togglebutton" },
     { name = "sphobjinv" },
+    { name = "tensorflow" },
 ]
 notebooks = [
+    { name = "ampform" },
     { name = "black" },
+    { name = "graphviz" },
     { name = "ipympl" },
+    { name = "jax" },
+    { name = "jaxlib" },
+    { name = "numba" },
+    { name = "numpy" },
     { name = "pandas" },
+    { name = "phasespace", extra = ["tf"] },
+    { name = "scipy" },
+    { name = "tensorflow" },
 ]
 numba = [
     { name = "numba" },
@@ -4404,6 +4431,7 @@ requires-dist = [
     { name = "tensorwaves", extras = ["doc"], marker = "extra == 'jupyter'" },
     { name = "tensorwaves", extras = ["jax"], marker = "extra == 'all'" },
     { name = "tensorwaves", extras = ["jax"], marker = "extra == 'types'" },
+    { name = "tensorwaves", extras = ["jax", "pwa"], marker = "extra == 'notebooks'" },
     { name = "tensorwaves", extras = ["jupyter"], marker = "extra == 'dev'" },
     { name = "tensorwaves", extras = ["notebooks"], marker = "extra == 'doc'" },
     { name = "tensorwaves", extras = ["numba"], marker = "extra == 'all'" },