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