def _perform_cache_optimization(
    saved_model_future: nodes.ValueNode,
    dataset_keys: Collection[analyzer_cache.DatasetKey],
    tensor_keys_to_paths: Dict[str, str],
    cache_dict: Optional[analyzer_cache.BeamAnalysisCache],
    num_phases: int,
) -> Tuple[
    Tuple[nodes.ValueNode],
    Optional[_IntermediateCacheType],
    Collection[nodes.ValueNode],
]:
  """Performs cache optimization on the given graph."""
  cache_output_nodes = {}
  optimize_visitor = _OptimizeVisitor(dataset_keys or {}, cache_dict,
                                      tensor_keys_to_paths, cache_output_nodes,
                                      num_phases)
  optimize_traverser = nodes.Traverser(optimize_visitor)
  optimized = optimize_traverser.visit_value_node(
      saved_model_future).flattened_view

  if cache_dict is None:
    assert not cache_output_nodes
    cache_output_nodes = None

  return (optimized, cache_output_nodes,
          optimize_visitor.get_detached_sideeffect_leafs())
Performs cache optimization on the given graph.
tensorflow_transform/beam/analysis_graph_builder.py