Update tests now that earliest supported TF version persists sparse tensor dense shapes

PiperOrigin-RevId: 546904332

Author: zoyahav

tensorflow_transform/beam/impl_test.py

@@ -4425,22 +4425,18 @@ def test3dSparseWithTFXIO(self):
             feature {
               name: "x$sparse_indices_0"
               type: INT
-              # TODO(b/184055743): Once TensorFlow is released with
-              # cl/514884437, uncomment.
-              # int_domain {
-              #   min: 0
-              #   max: 4
-              # }
+              int_domain {
+                min: 0
+                max: 4
+              }
             }
             feature {
               name: "x$sparse_indices_1"
               type: INT
-              # TODO(b/184055743): Once TensorFlow is released with
-              # cl/514884437 uncomment.
-              #  int_domain {
-              #   min: 0
-              #   max: 4
-              #  }
+               int_domain {
+                min: 0
+                max: 4
+               }
             }
             feature {
               name: "x$sparse_values"
@@ -4458,29 +4454,20 @@ def test3dSparseWithTFXIO(self):
               value_feature {
                 name: "x$sparse_values"
               }
-            }""", schema_pb2.Schema())
+            }""",
+            schema_pb2.Schema(),
+        )
         if not tft_unit.is_external_environment():
           expected_metadata.generate_legacy_feature_spec = False
 
-        # TODO(b/184055743): Once TensorFlow is released with cl/514884437,
-        # remove this.
-        def int_domain_cleared(schema):
-          result = schema_pb2.Schema()
-          result.CopyFrom(schema)
-          for f in result.feature:
-            f.ClearField('int_domain')
-          return result
-
-        self.assertProtoEquals(int_domain_cleared(transformed_metadata.schema),
-                               expected_metadata)
+        self.assertProtoEquals(transformed_metadata.schema, expected_metadata)
 
         beam_test_util.assert_that(
             transformed_data, self._MakeTransformOutputAssertFn(expected_data))
 
         def _assert_schemas_equal_fn(schema_dict_list):
           self.assertEqual(1, len(schema_dict_list))
-          self.assertProtoEquals(
-              int_domain_cleared(schema_dict_list[0].schema), expected_metadata)
+          self.assertProtoEquals(schema_dict_list[0].schema, expected_metadata)
 
         beam_test_util.assert_that(
             transformed_metadata.deferred_metadata,
@@ -4493,18 +4480,16 @@ def _assert_schemas_equal_fn(schema_dict_list):
       transformed_feature_spec = tft_out.transformed_feature_spec()
       self.assertLen(transformed_feature_spec, 1)
       self.assertIn('x', transformed_feature_spec)
-      self.assertIn(
+      self.assertEqual(
           transformed_feature_spec['x'],
-          (tf.io.SparseFeature(['x$sparse_indices_0', 'x$sparse_indices_1'],
-                               'x$sparse_values',
-                               tf.float32, [5, 5],
-                               already_sorted=True),
-           # TODO(b/184055743): Once TensorFlow is released with cl/514884437,
-           # remove this.
-           tf.io.SparseFeature(['x$sparse_indices_0', 'x$sparse_indices_1'],
-                               'x$sparse_values',
-                               tf.float32, [-1, -1],
-                               already_sorted=True)))
+          tf.io.SparseFeature(
+              ['x$sparse_indices_0', 'x$sparse_indices_1'],
+              'x$sparse_values',
+              tf.float32,
+              [5, 5],
+              already_sorted=True,
+          ),
+      )
 
       transformed_feature_spec['x'] = tf.io.SparseFeature(
           ['x$sparse_indices_0', 'x$sparse_indices_1'],

tensorflow_transform/impl_helper_test.py

@@ -725,10 +725,9 @@ def test_batched_placeholders_from_typespecs(self):
     self.assertEqual(features['dense_string'].dtype, tf.string)
 
     self.assertEqual(type(features['_sparse_underscored']), tf.SparseTensor)
-    # TODO(b/184055743): Once TensorFlow is released with cl/514884437,
-    # only keep the assertion of [None, None, 17].
-    self.assertIn(features['_sparse_underscored'].get_shape().as_list(),
-                  ([None, None, 17], [None, None, None]))
+    self.assertEqual(
+        features['_sparse_underscored'].get_shape().as_list(), [None, None, 17]
+    )
     self.assertEqual(features['_sparse_underscored'].dtype, tf.string)
 
     self.assertEqual(type(features['ragged_string']), tf.RaggedTensor)