Fixes to version parsing in tests.

PiperOrigin-RevId: 439402049

Author: varshaan

examples/census_example_v2_test.py

@@ -15,6 +15,7 @@
 
 import os
 import shutil
+from packaging import version
 
 import tensorflow.compat.v2 as tf
 import census_example_common
@@ -100,8 +101,8 @@ class CensusExampleV2Test(tft_test_case.TransformTestCase):
 
   def setUp(self):
     super().setUp()
-    if (not tf2.enabled() or
-        tft_test_case.is_external_environment() and tf.version.VERSION < '2.3'):
+    if (not tf2.enabled() or tft_test_case.is_external_environment() and
+        version.parse(tf.version.VERSION) < version.parse('2.3')):
       raise tft_test_case.SkipTest('This test requires TF version >= 2.3')
 
   def _get_data_dir(self):
@@ -113,7 +114,7 @@ def _get_working_dir(self):
         self._testMethodName)
 
   def _should_saved_model_load_work(self):
-    return tf.__version__ >= '2.2'
+    return version.parse(tf.__version__) >= version.parse('2.2')
 
   @tft_test_case.named_parameters([
       dict(

tensorflow_transform/impl_helper_test.py

@@ -17,6 +17,7 @@
 import os
 
 import numpy as np
+from packaging import version
 import pyarrow as pa
 import tensorflow as tf
 from tensorflow_transform import analyzers
@@ -659,7 +660,7 @@ def test_batched_placeholders_from_feature_spec(self):
                      [None, None])
     self.assertEqual(type(features['sparse_1d']), tf.SparseTensor)
     self.assertEqual(type(features['sparse_2d']), tf.SparseTensor)
-    if tf.__version__ >= '2':
+    if version.parse(tf.__version__) >= version.parse('2'):
       self.assertEqual(features['sparse_1d'].get_shape().as_list(), [None, 7])
       self.assertEqual(features['sparse_2d'].get_shape().as_list(),
                        [None, 2, 17])

tensorflow_transform/tf_utils.py

@@ -17,6 +17,7 @@
 import enum
 from typing import Callable, Optional, Tuple, Union
 
+from packaging import version
 import tensorflow as tf
 from tensorflow_transform import annotators
 from tensorflow_transform import common_types
@@ -543,7 +544,7 @@ def is_vocabulary_tfrecord_supported() -> bool:
     return False
   return ((hasattr(tf.data.experimental, 'DatasetInitializer') or
            hasattr(tf.lookup.experimental, 'DatasetInitializer')) and
-          tf.version.VERSION >= '2.4')
+          version.parse(tf.version.VERSION) >= version.parse('2.4'))
 
 
 # Used to decide which bucket boundary index to assign to a value.

tensorflow_transform/tf_utils_test.py

@@ -16,6 +16,7 @@
 import os
 
 import numpy as np
+from packaging import version
 import tensorflow as tf
 from tensorflow_transform import analyzers
 from tensorflow_transform import annotators
@@ -1583,9 +1584,10 @@ def _reduce_batch_count_mean_and_var_per_key(x, key):
               x=[[np.nan, np.nan, np.nan]],
               # Output of `tf.reduce_max` if all inputs are NaNs for older
               # versions of TF is -inf.
-              expected_x_minus_min=(-np.inf
-                                    if tf.__version__ < '2.4' else np.nan),
-              expected_x_max=-np.inf if tf.__version__ < '2.4' else np.nan,
+              expected_x_minus_min=(-np.inf if version.parse(tf.__version__) <
+                                    version.parse('2.4') else np.nan),
+              expected_x_max=(-np.inf if version.parse(tf.__version__) <
+                              version.parse('2.4') else np.nan),
               reduce_instance_dims=True,
               input_signature=[tf.TensorSpec([None, None], tf.float32)]),
           dict(