code: def get_traced_transform_fn(
      preprocessing_fn: Callable[
          [Mapping[str, common_types.TensorType]],
          Mapping[str, common_types.TensorType],
      ],
      input_signature: Mapping[str, tf.TypeSpec],
      tf_graph_context: graph_context.TFGraphContext,
      output_keys_to_name_map: Optional[Dict[str, str]] = None,
  ) -> tf.types.experimental.GenericFunction:
      """Get preprocessing_fn traced using tf.function.
  
      Args:
        preprocessing_fn: A user defined python function to be traced.
        input_signature: `tf.TypeSpec`s describing the inputs to the
          `preprocessing_fn`.
        tf_graph_context: A `TFGraphContext` context manager to invoke the
          `preprocessing_fn` in.
        output_keys_to_name_map: (Optional) A map from output dictionary keys to the
          names of the tensors that they represent.
  
      Returns:
        A tf.function object representing a function with the same input signature
        as `preprocessing_fn`.
        If `output_keys_to_name_map` is None or there are no more TFT analyzers to
        evaluate in the `preprocessing_fn`, the output signature of this
        tf.function
        is the same as the `preprocessing_fn`.
        Otherwise, its output signature contains the keys in
        `output_keys_to_name_map` and the tensor represented by the corresponding
        dictionary values.
      """
  
      assert all(
          [isinstance(s, tf.TypeSpec) for s in input_signature.values()]
      )
  
      # TODO(b/177672051): Investigate performance impact of enabling autograph.
      @tf.function(input_signature=[input_signature], autograph=False)
      def transform_fn(inputs):
          graph = ops.get_default_graph()
          # If any analyzers have already been evaluated, pass them using the
          # `graph_context.TFGraphContext`. This will be used in place of the analyzer
          # nodes.
          # The user defined `preprocessing_fn` may directly modify its inputs which
          # is not allowed in a tf.function. Hence, we make a copy here.
          inputs_copy = tf_utils.copy_tensors(inputs)
          with tf_graph_context:
              transformed_features = preprocessing_fn(inputs_copy)
          # An empty `TENSOR_REPLACEMENTS` collection symbolizes that there is no
          # analyzer left for Transform to evaluate. Either if this collection is
          # empty or if no specific outputs have been requested, return
          # the same output as `preprocessing_fn` (i.e, transformed_features).
          if (
              output_keys_to_name_map is None
              or not graph.get_collection(
                  analyzer_nodes.TENSOR_REPLACEMENTS
              )
          ):
              return transformed_features
          else:
              return {
                  key: graph.get_tensor_by_name(value)
                  for key, value in output_keys_to_name_map.items()
              }
  
      return transform_fn
  
imports: import functools
  import os
  import re
  from typing import Callable, Dict, List, Mapping, Optional, FrozenSet
  from absl import logging
  import numpy
  import pyarrow
  import tensorflow
  from tensorflow_transform import analyzer_nodes
  from tensorflow_transform import annotators
  from tensorflow_transform import common_types
  from tensorflow_transform import graph_context
  from tensorflow_transform import graph_tools
  from tensorflow_transform import schema_inference
  from tensorflow_transform import tf2_utils
  from tensorflow_transform import tf_utils
  from tensorflow_transform.output_wrapper import TFTransformOutput
  from tensorflow_transform.saved import saved_transform_io
  from tensorflow_transform.saved import saved_transform_io_v2
  from tensorflow_transform.tf_metadata import dataset_metadata
  from tensorflow_transform.tf_metadata import metadata_io
  from tensorflow_transform.tf_metadata import schema_utils
  from tfx_bsl.coders import example_coder
  from tfx_bsl.tfxio import tensor_to_arrow
  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 11 for calling a global function with example usage: get_traced_transform_fn(preprocessing_fn, input_signature, tf_graph_context, output_keys_to_name_map) and returns: transform_fn, transformed_features, unknown