diff --git a/var/spack/repos/builtin/packages/py-tensorflow-probability/package.py b/var/spack/repos/builtin/packages/py-tensorflow-probability/package.py
index 38af49c199..9f8360d432 100644
--- a/var/spack/repos/builtin/packages/py-tensorflow-probability/package.py
+++ b/var/spack/repos/builtin/packages/py-tensorflow-probability/package.py
@@ -23,8 +23,9 @@ class PyTensorflowProbability(Package):
 
     license("Apache-2.0")
 
-    # TODO: reactivate once TF 2.15 is ready https://github.com/spack/spack/pull/41069
-    # version("0.23.0", sha256="a00769550da9284acbd69e32a005507153ad39b0c190feca2bbbf6373366cc14")
+    # tf-keras does not (yet) exist on Spack, i.e. the TF extra won't work, see below
+    # version("0.24.0", sha256="3d75418cef09ea357ee879347133ab784fe4637a5ba251a8e06ef930dd970a3e")
+    version("0.23.0", sha256="a00769550da9284acbd69e32a005507153ad39b0c190feca2bbbf6373366cc14")
     version("0.22.1", sha256="9c1203b454aaeb48ac67dea862a411dba6b04f67c1e874e0e83bd1d7f13829a3")
     version("0.22.0", sha256="f9ce55b00c8069246d701c04eaafccde413355f6e76ccf9e549772ecfa0349a4")
     version("0.21.0", sha256="69b7510b38b2e48bcfb9ff570ef598d489e4f1bcbe13276f5dd91c878b8d56d1")
@@ -59,11 +60,13 @@ class PyTensorflowProbability(Package):
     # tensorflow_probability/python/__init__.py
     # TODO: reactivate the JAX versions once the JAX package is available with newer versions
     #  also add jaxlib as a dependency
-    # TODO: reactivate once TF 2.15 is ready https://github.com/spack/spack/pull/41069
 
     variant("py-tensorflow", default=False, description="Build with TensorFlow support")
     with when("+py-tensorflow"):
-        # depends_on("py-tensorflow@2.15", when="@0.23", type=("build", "run"))
+        # TODO(0.24): requires tf-keras on Spack
+        # depends_on("py-tf-keras", when="@0.24:", type=("build", "run"))
+        # depends_on("py-tensorflow@2.16", when="@0.24", type=("build", "run"))
+        depends_on("py-tensorflow@2.15", when="@0.23", type=("build", "run"))
         depends_on("py-tensorflow@2.14:2", when="@0.22", type=("build", "run"))
         depends_on("py-tensorflow@2.13:2", when="@0.21", type=("build", "run"))
         depends_on("py-tensorflow@2.12:2", when="@0.20", type=("build", "run"))
@@ -71,11 +74,11 @@ class PyTensorflowProbability(Package):
 
     # jaxlib is not required, as it's already a dependency of py-jax
     variant("py-jax", default=False, description="Build with JAX support")
-    with when("+py-jax"):  # TODO: reactivate once the JAX package is available with newer versions
-        # depends_on("py-jax@0.4.20:0.4", when="@0.23", type=("build", "run"))
-        # depends_on("py-jax@0.4.16:0.4", when="@0.22", type=("build", "run"))
-        # depends_on("py-jax@0.4.14:0.4", when="@0.21", type=("build", "run"))
-        # depends_on("py-jax@0.4.8:0.4", when="@0.20", type=("build", "run"))
+    with when("+py-jax"):
+        depends_on("py-jax@0.4.20:0.4", when="@0.23", type=("build", "run"))
+        depends_on("py-jax@0.4.16:0.4", when="@0.22", type=("build", "run"))
+        depends_on("py-jax@0.4.14:0.4", when="@0.21", type=("build", "run"))
+        depends_on("py-jax@0.4.8:0.4", when="@0.20", type=("build", "run"))
         depends_on("py-jax@0.3.25:3", when="@0.19", type=("build", "run"))
 
     depends_on(