def _get_output_to_inputs_map(
    output_signature: Mapping[str, common_types.TensorType]
) -> Dict[str, Iterable[tf.Tensor]]:
  """Gets all graph inputs that the tensors in output_signature depend on."""
  result = {}
  for name, output in output_signature.items():
    components = _get_component_tensors(output)
    sinks = [_as_operation(component) for component in components]
    # Ignore control dependencies when walking the graph as we only care about
    # which user defined inputs this output depends on.
    result[name] = graph_tools.retrieve_sources(
        sinks, ignore_control_dependencies=True)
  return result
Gets all graph inputs that the tensors in output_signature depend on.
tensorflow_transform/saved/saved_transform_io_v2.py