From fbaa55cc43cf2a0ca7c168f34039ee3c04fae1b5 Mon Sep 17 00:00:00 2001
From: Galen Andrew <galenandrew@google.com>
Date: Tue, 10 Aug 2021 13:23:10 -0700
Subject: [PATCH] Add dependency on tensorflow probability.

PiperOrigin-RevId: 389962756
---
 requirements.txt | 1 +
 setup.py         | 1 +
 2 files changed, 2 insertions(+)

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