import tensorflow as tf

      schema = schema_inference.infer_feature_schema(tensors, graph)

    expected_schema = schema_utils.schema_from_feature_spec(
        feature_spec, domains)
    self.assertEqual(schema, expected_schema)

  def test_infer_feature_schema_bad_rank(self):
    with tf.compat.v1.Graph().as_default() as graph:
      tensors = {
          'a': tf.compat.v1.placeholder(tf.float32, ()),
      }
    with self.assertRaises(ValueError):
      schema_inference.infer_feature_schema(tensors, graph)

  def test_bucketization_annotation(self):
    # TODO(b/132098015): Schema annotations aren't yet supported in OSS builds.
    # pylint: disable=g-import-not-at-top
    try:
