Automated rollback of commit 6c9ea7c

PiperOrigin-RevId: 514884437

Author: zoyahav

tensorflow_transform/beam/impl_test.py

@@ -4425,10 +4425,22 @@ def test3dSparseWithTFXIO(self):
             feature {
               name: "x$sparse_indices_0"
               type: INT
+              # TODO(b/184055743): Once TensorFlow is released with
+              # cl/342914534, uncomment.
+              # int_domain {
+              #   min: 0
+              #   max: 4
+              # }
             }
             feature {
               name: "x$sparse_indices_1"
               type: INT
+              # TODO(b/184055743): Once TensorFlow is released with
+              # cl/342914534 uncomment.
+              #  int_domain {
+              #   min: 0
+              #   max: 4
+              #  }
             }
             feature {
               name: "x$sparse_values"
@@ -4450,14 +4462,28 @@ def test3dSparseWithTFXIO(self):
         if not tft_unit.is_external_environment():
           expected_metadata.generate_legacy_feature_spec = False
 
-        self.assertProtoEquals(transformed_metadata.schema, expected_metadata)
+        # TODO(b/184055743): Once TensorFlow is released with cl/342914534,
+        # remove this.
+        # TODO(b/184057384): Even with cl/342914534,
+        # transformed_metadata.deferred_schema still does not contain the shape
+        # information about the SparseTensor.
+        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)
 
         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(schema_dict_list[0].schema, expected_metadata)
+          self.assertProtoEquals(
+              int_domain_cleared(schema_dict_list[0].schema), expected_metadata)
 
         beam_test_util.assert_that(
             transformed_metadata.deferred_metadata,
@@ -4468,15 +4494,26 @@ def _assert_schemas_equal_fn(schema_dict_list):
       dataset = tf.data.TFRecordDataset(materialize_path)
       tft_out = tft.TFTransformOutput(transform_output_path)
       transformed_feature_spec = tft_out.transformed_feature_spec()
-      self.assertEqual(
-          transformed_feature_spec, {
-              'x':
-                  tf.io.SparseFeature(
-                      ['x$sparse_indices_0', 'x$sparse_indices_1'],
-                      'x$sparse_values',
-                      tf.float32, [-1, -1],
-                      already_sorted=True)
-          })
+      self.assertLen(transformed_feature_spec, 1)
+      self.assertIn('x', transformed_feature_spec)
+      self.assertIn(
+          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/342914534,
+           # remove this.
+           # TODO(b/184057384): Even with cl/342914534, still may not contain
+           # the shape information about the SparseTensor.
+           tf.io.SparseFeature(['x$sparse_indices_0', 'x$sparse_indices_1'],
+                               'x$sparse_values',
+                               tf.float32, [-1, -1],
+                               already_sorted=True)))
+
+      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)
 
       def parse_fn(serialized_input):
         result = tf.io.parse_single_example(serialized_input,
@@ -4493,7 +4530,7 @@ def parse_fn(serialized_input):
     expected_sparse_components = [
         np.array([[arr] for arr in zip(x_idx0, x_idx1)]),
         np.array([[x] for x in x_data]),
-        np.array([[-1, -1]] * len(x_data))
+        np.array([[5, 5]] * len(x_data))
     ]
     self.assertLen(transformed_sparse_components,
                    len(expected_sparse_components))

tensorflow_transform/impl_helper_test.py

@@ -725,10 +725,10 @@ def test_batched_placeholders_from_typespecs(self):
     self.assertEqual(features['dense_string'].dtype, tf.string)
 
     self.assertEqual(type(features['_sparse_underscored']), tf.SparseTensor)
-    # TODO(zoyahav): Change last dimension size to 17 once SparseTensors propagate
-    # static dense_shape from typespec correctly.
-    self.assertEqual(features['_sparse_underscored'].get_shape().as_list(),
-                     [None, None, None])
+    # TODO(b/184055743): Once TensorFlow is released with cl/342914534,
+    # 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'].dtype, tf.string)
 
     self.assertEqual(type(features['ragged_string']), tf.RaggedTensor)