def perform_combiner_packing_optimization(saved_model_future,
                                          cache_value_nodes, num_phases):
  """Optimizes the graph by packing possible combine nodes."""
  # Inspect the graph to identify all the packable combines.
  inspect_acc_combine_visitor = _InspectAccumulateCombineVisitor()
  inspect_acc_combine_traverser = nodes.Traverser(inspect_acc_combine_visitor)
  _ = inspect_acc_combine_traverser.visit_value_node(saved_model_future)

  packable_combines = inspect_acc_combine_visitor.packable_combines
  # Do not pack if we have only a single combine in the group.
  packable_combines = {
      label: group for label, group in packable_combines.items()
      if len(group) > 1
  }

  pack_acc_combine_visitor = _PackAccumulateCombineVisitor(packable_combines)
  pack_acc_combine_traverser = nodes.Traverser(pack_acc_combine_visitor)
  saved_model_future = pack_acc_combine_traverser.visit_value_node(
      saved_model_future)

  # Replace cache nodes to point to the corresponding new nodes.
  cache_value_nodes = _update_cache_value_node_references(
      cache_value_nodes, pack_acc_combine_traverser)

  # TODO(b/134414978): Consider also packing the merges even when we have
  # multiple phases.
  if num_phases > 1:
    return (saved_model_future, cache_value_nodes)

  # Identify the merge combines that can be packed together.
  inspect_merge_combine_visitor = _InspectMergeCombineVisitor()
  inspect_merge_combine_traverser = nodes.Traverser(
      inspect_merge_combine_visitor)
  _ = inspect_merge_combine_traverser.visit_value_node(saved_model_future)

  # Only pack if we have more than one merge combines.
  if len(inspect_merge_combine_visitor.packable_combine_extract_outputs) <= 1:
    return (saved_model_future, cache_value_nodes)

  # Add flatten and packed merge nodes.
  pack_merge_combine_visitor = _PackMergeCombineVisitor(
      packable_combine_extract_outputs=
      inspect_merge_combine_visitor.packable_combine_extract_outputs)
  pack_merge_combine_traverser = nodes.Traverser(pack_merge_combine_visitor)
  saved_model_future = pack_merge_combine_traverser.visit_value_node(
      saved_model_future)
  # Replace cache nodes to point to the corresponding new nodes.
  cache_value_nodes = _update_cache_value_node_references(
      cache_value_nodes, pack_merge_combine_traverser)

  # Remove redundant flatten and packed merge nodes.
  remove_redundant_visitor = _RemoveRedundantPackedMergeCombineVisitor(
      final_packed_merge_combine_label=
      pack_merge_combine_visitor.final_packed_merge_combine_label)
  remove_redundant_traverser = nodes.Traverser(remove_redundant_visitor)
  saved_model_future = remove_redundant_traverser.visit_value_node(
      saved_model_future)
  # Replace cache nodes to point to the corresponding new nodes.
  cache_value_nodes = _update_cache_value_node_references(
      cache_value_nodes, remove_redundant_traverser)

  return (saved_model_future, cache_value_nodes)
Optimizes the graph by packing possible combine nodes.
tensorflow_transform/beam/combiner_packing_util.py