import tensorflow as tf


      # 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)
      # Create a session to actually evaluate the annotations and extract the
      # the output schema with annotations applied.
      with tf.compat.v1.Session(graph=graph) as session:
        schema = schema_inference.infer_feature_schema(outputs, graph, session)
        self.assertLen(schema.feature, 2)
        for feature in schema.feature:
          self.assertLen(feature.annotation.extra_metadata, 1)
          for annotation in feature.annotation.extra_metadata:

            # Extract the annotated message and validate its contents
            message = annotations_pb2.BucketBoundaries()
