def _get_schema_overrides(graph,
                          tensor_name_to_key_map,
                          tensor_collection_key,
                          overrides_keys,
                          default_tensor_name=None):
  """Obtain schema overrides from graph collections.

  For every tensor in the `tensor_collection_key` collection, the corresponding
  feature name is in `tensor_name_to_key_map` and its schema overrides are in
  the graph collections defined by keys in `overrides_keys`.
  If a tensor does not exist in `tensor_name_to_key_map` but its name starts
  with `default_tensor_name` (if provided), the overrides are returned with this
  key.

  Args:
    graph: A `FuncGraph`.
    tensor_name_to_key_map: A dictionary from tensor name to output feature key.
    tensor_collection_key: Key for the graph collection that contains list of
      tensors to annotate.
    overrides_keys: A list of graph collection keys that contain schema
      overrides/annotations.
    default_tensor_name: (Optional) A String. If provided, use as feature key if
      a tensor in the graph collections is not in `tensor_name_to_key_map`.

  Returns:
    A dictionary from graph collection keys to lists of features and their
    schema overrides/annotations.

  """
  tensors = graph.get_collection(tensor_collection_key)
  overrides_list = [graph.get_collection(k) for k in overrides_keys]

  result = collections.defaultdict(list)
  if len(tensors) != len(overrides_list[0]) or any(
      len(lst) != len(overrides_list[0]) for lst in overrides_list
  ):
    raise ValueError(
        f'Unexpected collections lengths. tensors: {tensors}, overrides_list:'
        f' {overrides_list}'
    )
  for tensor, overrides_tuple in zip(tensors, zip(*overrides_list)):
    if tensor.name in tensor_name_to_key_map:
      result[tensor_collection_key].append(tensor_name_to_key_map[tensor.name])
    else:
      if default_tensor_name is None:
        continue
      tensor_name = tensor.name.split('/')[-1]
      if tensor.dtype == tf.string and tensor_name.startswith(
          default_tensor_name):
        result[tensor_collection_key].append(default_tensor_name)
      else:
        continue

    # If a feature name was added to the result list for tensor_collection_key,
    # add its annotations as well.
    assert len(overrides_keys) == len(overrides_tuple)
    for overrides_key, override in zip(overrides_keys, overrides_tuple):
      result[overrides_key].append(override)
  return result
Obtain schema overrides from graph collections.

  For every tensor in the `tensor_collection_key` collection, the corresponding
  feature name is in `tensor_name_to_key_map` and its schema overrides are in
  the graph collections defined by keys in `overrides_keys`.
  If a tensor does not exist in `tensor_name_to_key_map` but its name starts
  with `default_tensor_name` (if provided), the overrides are returned with this
  key.

  Args:
    graph: A `FuncGraph`.
    tensor_name_to_key_map: A dictionary from tensor name to output feature key.
    tensor_collection_key: Key for the graph collection that contains list of
      tensors to annotate.
    overrides_keys: A list of graph collection keys that contain schema
      overrides/annotations.
    default_tensor_name: (Optional) A String. If provided, use as feature key if
      a tensor in the graph collections is not in `tensor_name_to_key_map`.

  Returns:
    A dictionary from graph collection keys to lists of features and their
    schema overrides/annotations.

  
tensorflow_transform/schema_inference.py