diff --git a/requirements.txt b/requirements.txt
index 955d819..7994d55 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,3 +4,4 @@ tensorflow-estimator>=2.3.0
 attrs>=21.2.0
 mpmath
 dm-tree~=0.1.1
+tensorflow-probability>=0.13.0
diff --git a/setup.py b/setup.py
index 155802d..4ca4064 100644
--- a/setup.py
+++ b/setup.py
@@ -26,6 +26,7 @@ setup(
         'attrs>=21.2.0',  # for tree_aggregation_query.py.
         'mpmath',  # used in tests only
         'dm-tree~=0.1.1',  # used in tests only
+        'tensorflow-probability>=0.13.0',  # For discrete Gaussian.
     ],
     # Explicit dependence on TensorFlow is not supported.
     # See https://github.com/tensorflow/tensorflow/issues/7166