code: def test_bucketization_annotation(self, use_compat_v1):
      if not use_compat_v1:
          test_case.skip_if_not_tf2("Tensorflow 2.x required")
  
      def preprocessing_fn(_):
          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([0.5, 1.5]), axis=0
          )
          boundaries_bar = tf.expand_dims(
              tf.convert_to_tensor([0.1, 0.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
          )
          return outputs
  
      schema = self._get_schema(
          preprocessing_fn, use_compat_v1, create_session=True
      )
      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()
              annotation.Unpack(message)
              if feature.name == "Bucketized_foo":
                  self.assertAllClose(
                      list(message.boundaries), [0.5, 1.5]
                  )
              elif feature.name == "Bucketized_bar":
                  self.assertAllClose(
                      list(message.boundaries), [0.1, 0.2]
                  )
              else:
                  raise RuntimeError("Unexpected features in schema")
  
imports: import functools
  import os
  import tensorflow
  from tensorflow_transform import analyzers
  from tensorflow_transform import common
  from tensorflow_transform import graph_context
  from tensorflow_transform import mappers
  from tensorflow_transform import schema_inference
  from tensorflow_transform import tf2_utils
  from tensorflow_transform import test_case
  from tensorflow_transform.tf_metadata import schema_utils_legacy
  from tensorflow_transform.tf_metadata import schema_utils
  from google.protobuf import text_format
  import unittest
  from tensorflow_metadata.proto.v0 import schema_pb2
label_desc: Use image node_id 5 for calling the SchemaInferenceTest obj's underlying member method code with example usage: obj.test_bucketization_annotation(use_compat_v1) and returns: outputs