From 5aeb91aff2a1fd8f1d39e181dfbc27d6206a7701 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 21 Nov 2024 07:04:54 -0800
Subject: [PATCH] Replace distutils's LooseVersion with packaging's Version

PiperOrigin-RevId: 698767136
---
 tensorflow_privacy/privacy/bolt_on/__init__.py          | 6 +++---
 tensorflow_privacy/privacy/dp_query/tree_range_query.py | 6 +++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensorflow_privacy/privacy/bolt_on/__init__.py b/tensorflow_privacy/privacy/bolt_on/__init__.py
index 4271f08b..382801d2 100644
--- a/tensorflow_privacy/privacy/bolt_on/__init__.py
+++ b/tensorflow_privacy/privacy/bolt_on/__init__.py
@@ -13,16 +13,16 @@
 # limitations under the License.
 """BoltOn Method for privacy."""
 
-from distutils import version
 import sys
 
+from packaging import version
 import tensorflow.compat.v1 as tf
 
-if version.LooseVersion(tf.__version__) < version.LooseVersion("2.0.0"):
+if version.Version(tf.__version__) < version.Version("2.0.0"):
   raise ImportError("Please upgrade your version "
                     "of tensorflow from: {0} to at least 2.0.0 to "
                     "use privacy/bolt_on".format(
-                        version.LooseVersion(tf.__version__)))
+                        version.Version(tf.__version__)))
 if hasattr(sys, "skip_tf_privacy_import"):  # Useful for standalone scripts.
   pass
 else:
diff --git a/tensorflow_privacy/privacy/dp_query/tree_range_query.py b/tensorflow_privacy/privacy/dp_query/tree_range_query.py
index 9cfee022..8213fa41 100644
--- a/tensorflow_privacy/privacy/dp_query/tree_range_query.py
+++ b/tensorflow_privacy/privacy/dp_query/tree_range_query.py
@@ -16,11 +16,11 @@
 'Offline' means all the leaf nodes are ready before the protocol starts.
 """
 
-import distutils
 import math
 from typing import Any, NamedTuple, Optional
 
 import dp_accounting
+import packaging
 import tensorflow as tf
 from tensorflow_privacy.privacy.dp_query import distributed_discrete_gaussian_query
 from tensorflow_privacy.privacy.dp_query import dp_query
@@ -270,8 +270,8 @@ def build_distributed_discrete_gaussian_query(cls,
 
 def _get_add_noise(stddev, seed: Optional[int] = None):
   """Utility function to decide which `add_noise` to use according to tf version."""
-  if distutils.version.LooseVersion(
-      tf.__version__) < distutils.version.LooseVersion('2.0.0'):
+  if packaging.version.Version(
+      tf.__version__) < packaging.version.Version('2.0.0'):
 
     # The seed should be only used for testing purpose.
     if seed is not None: