code: def get_traced_metadata_fn(
      preprocessing_fn: Callable[
          [Mapping[str, common_types.TensorType]],
          Mapping[str, common_types.TensorType],
      ],
      structured_inputs: Mapping[str, common_types.TensorType],
      tf_graph_context: graph_context.TFGraphContext,
      evaluate_schema_overrides: bool,
  ) -> tf.types.experimental.GenericFunction:
      """Get a tf.function that returns a dictionary of annotations.
  
      Annotations are added to graph collections keyed by graph tensor names when
      `preprocessing_fn` is being traced. The metadata fn defined by this method
      converts the graph tensor names to output feature keys.
  
      If `evaluate_schema_overrides` is True, tracing the `preprocessing_fn` will
      add overrides for feature ranges (min/max) and/or feature protos to the graph
      collection, if applicable. These overrides are returned when the function
      returned by this method is invoked.
  
      Args:
        preprocessing_fn: A user defined python function to be traced.
        structured_inputs: A dictionary of placeholder inputs to `preprocessing_fn`.
        tf_graph_context: A `TFGraphContext` context manager to invoke the
          `preprocessing_fn` in.
        evaluate_schema_overrides: If `False`, the returned dictionary contains a
          single key `_TF_METADATA_TENSOR_COLLECTION` as all other annotations are
          deferred. Else, the returned dictionary contains several deferred
          annotations.
  
      Returns:
        A tf.function which when invoked returns a dictionary whose keys represent
        the types of annotations and the values are collections of feature
        keys/annotations.
      """
  
      # Since this is a TFT-internal function with constant outputs, autograph will
      # not affect its behavior. It will only increase tracing time, if enabled.
      # Hence, trace with `autograph=False` here.
      @tf.function(input_signature=[], autograph=False)
      def metadata_fn():
          graph = ops.get_default_graph()
          inputs = tf2_utils.supply_missing_inputs(
              structured_inputs, batch_size=1
          )
          with tf_graph_context:
              transformed_features = preprocessing_fn(inputs)
  
          # Get a map from tensor value names to feature keys.
          reversed_features = _get_tensor_value_to_key_map(
              transformed_features
          )
  
          result = collections.defaultdict(list)
          if not evaluate_schema_overrides:
              schema_override_tensors = graph.get_collection(
                  _TF_METADATA_TENSOR_COLLECTION
              )
              for tensor in schema_override_tensors:
                  if tensor.name in reversed_features:
                      result[_TF_METADATA_TENSOR_COLLECTION].append(
                          reversed_features[tensor.name]
                      )
          else:
              # Obtain schema overrides for feature tensor ranges.
              result.update(
                  _get_schema_overrides(
                      graph,
                      reversed_features,
                      _TF_METADATA_TENSOR_COLLECTION,
                      [
                          _TF_METADATA_TENSOR_MIN_COLLECTION,
                          _TF_METADATA_TENSOR_MAX_COLLECTION,
                      ],
                  )
              )
              # Obtain schema overrides for feature protos. If no feature tensor is in
              # the `_TF_METADATA_EXTRA_ANNOTATION` collection for a specified
              # annotation, `_TF_METADATA_EXTRA_ANNOTATION_GLOBAL` is used as the
              # feature name to indicate that this annotation should be added to the
              # global schema.
              result.update(
                  _get_schema_overrides(
                      graph,
                      reversed_features,
                      _TF_METADATA_EXTRA_ANNOTATION,
                      [
                          _TF_METADATA_EXTRA_ANNOTATION_TYPE_URL,
                          _TF_METADATA_EXTRA_ANNOTATION_PROTO,
                      ],
                      _TF_METADATA_EXTRA_ANNOTATION_GLOBAL,
                  )
              )
              result[_METADATA_SPARSE_OUTPUT_OVERRIDES_FIELD] = [
                  _get_sparse_output_annotations_v2(
                      graph, reversed_features
                  )
              ]
          return result
  
      # Though this tf.function is not serialized to SavedModel, if it is capturing
      # any resources, they need to be tracked to ensure they are not garbage
      # collected.
      # We strip control dependencies from this function as we only want to evaluate
      # annotations which are constants available at graph construction time and
      # have no dependency on inputs.
      # TODO(b/149352022): Re-factor metadata computation when it is possible to
      # evaluate non output tensors from a func graph.
      module = tf_graph_context.module_to_export
      saved_transform_io_v2.trace_and_update_module(
          module,
          metadata_fn,
          "metadata_fn",
          strip_control_dependencies=True,
      )
      return module.metadata_fn
  
imports: import collections
  import itertools
  from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Union
  from absl import logging
  import tensorflow
  from tensorflow_transform import common
  from tensorflow_transform import common_types
  from tensorflow_transform import graph_context
  from tensorflow_transform import tf2_utils
  from tensorflow_transform import tf_utils
  from tensorflow_transform.saved import saved_transform_io_v2
  from tensorflow_transform.tf_metadata import schema_utils
  from tfx_bsl.tfxio import tensor_representation_util
  from google.protobuf import any_pb2
  from tensorflow.python.eager import function
  from tensorflow.python.framework import ops
  from tensorflow_metadata.proto.v0 import schema_pb2
label_desc: Use image node_id 18 for calling a global function with example usage: get_traced_metadata_fn(preprocessing_fn, structured_inputs, tf_graph_context, evaluate_schema_overrides) and returns: module, result