From f99a74c7a4c6be2c07d761183e92e624ae6fef46 Mon Sep 17 00:00:00 2001
From: Steve Chien <schien@google.com>
Date: Tue, 20 Dec 2022 11:48:44 -0800
Subject: [PATCH] Fix dependencies required by `privacy_tests`. Update
 `distutils` to `packaging`.

PiperOrigin-RevId: 496713867
---
 requirements.txt                                      | 3 +++
 setup.py                                              | 4 ++++
 tensorflow_privacy/privacy/dp_query/gaussian_query.py | 7 +++----
 tensorflow_privacy/privacy/privacy_tests/BUILD        | 9 +--------
 4 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 3bcbd66..50c009c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -24,12 +24,15 @@ absl-py>=1.0,==1.*
 attrs~=21.4
 dm-tree==0.1.7
 dp-accounting==0.3.0
+immutabledict~=2.2
 matplotlib~=3.3
 mpmath~=1.2
 numpy~=1.21
+packaging~=22.0
 pandas~=1.4
 scikit-learn>=1.0,==1.*
 scipy~=1.7
+statsmodels~=0.13
 tensorflow-datasets~=4.5
 tensorflow-estimator~=2.4
 tensorflow-probability==0.15.0
diff --git a/setup.py b/setup.py
index fbaccc9..ab83cc8 100644
--- a/setup.py
+++ b/setup.py
@@ -31,11 +31,15 @@ setup(
         'attrs~=21.4',
         'dm-tree==0.1.7',
         'dp-accounting==0.3.0',
+        'immutabledict~=2.2',
         'matplotlib~=3.3',
         'numpy~=1.21',
+        'packaging~=22.0',
         'pandas~=1.4',
+        'parameterized~=0.8',
         'scikit-learn>=1.0,==1.*',
         'scipy~=1.7',
+        'statsmodels~=0.13',
         'tensorflow-datasets~=4.5',
         'tensorflow-estimator~=2.4',
         'tensorflow-probability==0.15.0',
diff --git a/tensorflow_privacy/privacy/dp_query/gaussian_query.py b/tensorflow_privacy/privacy/dp_query/gaussian_query.py
index a4aeeca..b50fae6 100644
--- a/tensorflow_privacy/privacy/dp_query/gaussian_query.py
+++ b/tensorflow_privacy/privacy/dp_query/gaussian_query.py
@@ -14,9 +14,8 @@
 """Implements DPQuery interface for Gaussian sum queries."""
 
 import collections
-import distutils
-
 import dp_accounting
+import packaging
 import tensorflow as tf
 from tensorflow_privacy.privacy.dp_query import dp_query
 
@@ -79,8 +78,8 @@ class GaussianSumQuery(dp_query.SumAggregationDPQuery):
 
   def get_noised_result(self, sample_state, global_state):
     """Implements `tensorflow_privacy.DPQuery.get_noised_result`."""
-    if distutils.version.LooseVersion(
-        tf.__version__) < distutils.version.LooseVersion('2.0.0'):
+    if packaging.version.Version(
+        tf.__version__) < packaging.version.Version('2.0.0'):
 
       def add_noise(v):
         return v + tf.random.normal(
diff --git a/tensorflow_privacy/privacy/privacy_tests/BUILD b/tensorflow_privacy/privacy/privacy_tests/BUILD
index ad133d3..af20a7c 100644
--- a/tensorflow_privacy/privacy/privacy_tests/BUILD
+++ b/tensorflow_privacy/privacy/privacy_tests/BUILD
@@ -15,10 +15,7 @@ py_test(
     srcs = ["utils_test.py"],
     python_version = "PY3",
     srcs_version = "PY3",
-    deps = [
-        ":utils",
-        "//third_party/py/parameterized",
-    ],
+    deps = [":utils"],
 )
 
 py_test(
@@ -36,8 +33,4 @@ py_library(
 py_library(
     name = "epsilon_lower_bound",
     srcs = ["epsilon_lower_bound.py"],
-    deps = [
-        "//third_party/py/immutabledict",
-        "//third_party/py/statsmodels",
-    ],
 )