Suppress distutils deprecation warning from distutils Version classes.

PiperOrigin-RevId: 795983785

Author: Googler

tensorflow_probability/python/BUILD

@@ -35,6 +35,7 @@ py_library(
     deps = [
         ":version",
         # distutils dep,
+        # packaging dep,
         "//tensorflow_probability/python/bijectors",
         "//tensorflow_probability/python/debugging",
         "//tensorflow_probability/python/distributions",

tensorflow_probability/python/__init__.py

@@ -47,15 +47,23 @@ def _validate_tf_environment(package):
     raise
 
   import distutils.version
+  import packaging.version
+
+  def _loose_version(version: str):
+    try:
+      return packaging.version.Version(version)
+    except packaging.version.InvalidVersion:
+      return distutils.version.LooseVersion(version)
 
   #
   # Update this whenever we need to depend on a newer TensorFlow release.
   #
   required_tensorflow_version = '2.18'
 #   required_tensorflow_version = '1.15'  # Needed internally -- DisableOnExport
 
-  if (distutils.version.LooseVersion(tf.__version__) <
-      distutils.version.LooseVersion(required_tensorflow_version)):
+  if _loose_version(tf.__version__) < _loose_version(
+      required_tensorflow_version
+  ):
     raise ImportError(
         'This version of TensorFlow Probability requires TensorFlow '
         'version >= {required}; Detected an installation of version {present}. '