diff --git a/pkgs/development/python-modules/tensorflow-probability/default.nix b/pkgs/development/python-modules/tensorflow-probability/default.nix
index 1daaac6ae7eb..760e9a0e9d1c 100644
--- a/pkgs/development/python-modules/tensorflow-probability/default.nix
+++ b/pkgs/development/python-modules/tensorflow-probability/default.nix
@@ -1,7 +1,5 @@
 { lib
 , fetchFromGitHub
-, fetchpatch
-, bazel_0_26
 , buildBazelPackage
 , buildPythonPackage
 , python
@@ -11,47 +9,33 @@
 , tensorflow
 , six
 , numpy
+, dm-tree
+, keras
 , decorator
 , cloudpickle
 , gast
 , hypothesis
 , scipy
+, pandas
+, mpmath
 , matplotlib
 , mock
 , pytest
 }:
 
 let
-  version = "0.8.0";
+  version = "0.15.0";
   pname = "tensorflow_probability";
 
   # first build all binaries and generate setup.py using bazel
   bazel-wheel = buildBazelPackage {
-    bazel = bazel_0_26;
-
     name = "${pname}-${version}-py2.py3-none-any.whl";
-
     src = fetchFromGitHub {
       owner = "tensorflow";
       repo = "probability";
-      rev = version;
-      sha256 = "07cm8zba8n0ihzdm3k4a4rsg5v62xxsfvcw4h0niz91c0parqjqy";
+      rev = "v" + version;
+      sha256 = "155fgmra90s08vjnp61qxdrpzq74xa3kdzhgdkavwgc25pvxn3mi";
     };
-
-    patches = [
-      (fetchpatch {
-        name = "gast-0.3.patch";
-        url = "https://github.com/tensorflow/probability/commit/ae7a9d9771771ec1e7755a3588b9325f050a84cc.patch";
-        sha256 = "0kfhx30gshm8f3945na9yjjik71r20qmjzifbigaj4l8dwd9dz1a";
-        excludes = ["testing/*"];
-      })
-      (fetchpatch {
-        name = "cloudpickle-1.2.patch";
-        url = "https://github.com/tensorflow/probability/commit/78ef12b5afe3f567d16c70b74015ed1ddff1b0c8.patch";
-        sha256 = "12ms2xcljvvrnig0j78s3wfv4yf3bm5ps4rgfgv5lg2a8mzpc1ga";
-      })
-    ];
-
     nativeBuildInputs = [
       # needed to create the output wheel in installPhase
       python
@@ -64,7 +48,7 @@ let
     bazelTarget = ":pip_pkg";
 
     fetchAttrs = {
-      sha256 = "1qw7vkwnxy45z4vm94isq5m96xiz35sigag7vjg1xb2sklbymxh8";
+      sha256 = "0sgxdlw5x3dydy53l10vbrj8smh78b7r1wff8jxcgp4w69mk8zfm";
     };
 
     buildAttrs = {
@@ -98,34 +82,31 @@ in buildPythonPackage {
     decorator
     cloudpickle
     gast
+    dm-tree
+    keras
   ];
 
   # Listed here:
-  # https://github.com/tensorflow/probability/blob/f01d27a6f256430f03b14beb14d37def726cb257/testing/run_tests.sh#L58
+  # https://github.com/tensorflow/probability/blob/f3777158691787d3658b5e80883fe1a933d48989/testing/dependency_install_lib.sh#L83
   checkInputs = [
     hypothesis
     pytest
     scipy
+    pandas
+    mpmath
     matplotlib
     mock
   ];
 
-  # actual checks currently fail because for some reason
-  # tf.enable_eager_execution is called too late. Probably because upstream
-  # intents these tests to be run by bazel, not plain pytest.
-  # checkPhase = ''
-  #   # tests need to import from other test files
-  #   export PYTHONPATH="$PWD/tensorflow-probability:$PYTHONPATH"
-  #   py.test
-  # '';
+  # Ideally, we run unit tests with pytest, but in checkPhase, only the Bazel-build wheel is available.
+  # But it seems not guaranteed that running the tests with pytest will even work, see
+  # https://github.com/tensorflow/probability/blob/c2a10877feb2c4c06a4dc58281e69c37a11315b9/CONTRIBUTING.md?plain=1#L69
+  # Ideally, tests would be run using Bazel. For now, lets's do a...
 
   # sanity check
-  checkPhase = ''
-    python -c 'import tensorflow_probability'
-  '';
+  pythonImportsCheck = [ "tensorflow_probability" ];
 
   meta = with lib; {
-    broken = true;  # tf-probability 0.8.0 is not compatible with tensorflow 2.3.2
     description = "Library for probabilistic reasoning and statistical analysis";
     homepage = "https://www.tensorflow.org/probability/";
     license = licenses.asl20;