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
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.
  
tensorflow_transform/impl_helper.py