import tensorflow as tf

    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:
      from tensorflow_transform import annotations_pb2
    except ImportError:
      return
    # pylint: enable=g-import-not-at-top
    with tf.compat.v1.Graph().as_default() as graph:
      inputs = {
          'foo': tf.convert_to_tensor([0, 1, 2, 3]),
          'bar': tf.convert_to_tensor([0, 2, 0, 2]),
      }
      boundaries_foo = tf.expand_dims(tf.convert_to_tensor([.5, 1.5]), axis=0)
      boundaries_bar = tf.expand_dims(tf.convert_to_tensor([.1, .2]), axis=0)
      outputs = {}

      # tft.apply_buckets will annotate the feature in the output schema to
      # indicate the bucket boundaries that were applied.
      outputs['Bucketized_foo'] = mappers.apply_buckets(inputs['foo'],
                                                        boundaries_foo)
      outputs['Bucketized_bar'] = mappers.apply_buckets(inputs['bar'],
                                                        boundaries_bar)
