import tensorflow as tf

      outputs = {
          'foo': tf.convert_to_tensor([0, 1, 2, 3], dtype=tf.int64),
          'bar': tf.convert_to_tensor([0, 2, 0, 2], dtype=tf.int64),
      }

      # Annotate an arbitrary proto at the schema level (not sure what global
      # schema boundaries would mean, but hey I'm just a test).
      boundaries = tf.constant([[1.0]])
      message_type = annotations_pb2.BucketBoundaries.DESCRIPTOR.full_name
      sizes = tf.expand_dims([tf.size(boundaries)], axis=0)
      message_proto = tf.raw_ops.EncodeProto(
          sizes=sizes, values=[tf.cast(boundaries, tf.float32)],
          field_names=['boundaries'], message_type=message_type)[0]
      type_url = os.path.join('type.googleapis.com', message_type)
      schema_inference.annotate(type_url, message_proto)

      with tf.compat.v1.Session(graph=graph) as session:
        schema = schema_inference.infer_feature_schema(outputs, graph, session)
        self.assertLen(schema.annotation.extra_metadata, 1)
        for annotation in schema.annotation.extra_metadata:
          # Extract the annotated message and validate its contents
          message = annotations_pb2.BucketBoundaries()
