class_docstr: Used to perform in-graph replication on a computational graph.
  
  That is, Replicator is used to create _n_ replicas of a `tf.Graph`,
  with all replicas being stored in a new `tf.Graph` and delineated
  with name scopes.
  
  This also handles ControlFlowContexts, updating variable and gradient
  information, etc.
class_docstr_tok: ["Used","to","perform","in-graph","replication","on","a","computational","graph",".","That","is",",","Replicator","is","used","to","create","_n_","replicas","of","a","`","tf.Graph","`",",","with","all","replicas","being","stored","in","a","new","`","tf.Graph","`","and","delineated","with","name","scopes",".","This","also","handles","ControlFlowContexts",",","updating","variable","and","gradient","information",",","etc","."]
docstr: Replicate the entire graph as many times as num_replica.
  
  Args:
      graph_item: the original graph item
  
  Returns: The new graph item
docstr_tok: ["Replicate","the","entire","graph","as","many","times","as","num_replica",".","Args",":","graph_item",":","the","original","graph","item","Returns",":","The","new","graph","item"]
code: def replicate(self, graph_item):
      """
      Replicate the entire graph as many times as num_replica.
  
      Args:
          graph_item: the original graph item
  
      Returns: The new graph item
      """
      item = GraphItem(graph=ops.Graph())
      fwd_ctx, bwd_ctx = self._collect_while_context(graph_item.graph)
      with item.graph.as_default():
          gdef = graph_item.graph.as_graph_def()
          for i in range(self._num_local_replicas):
              # Replicate ops
              with ops.device(
                  self._replica_device_placer(replica_id=i)
              ):
                  import_graph_def(gdef, name=replica_prefix(i))
  
              # Replicate while_loop context (control_flow) if needed.
              # The order matters -- We must replicate bwd context first, then forward context.
              # TODO(Zeya): To handle cases when there are nested while loops, in which we must replicate
              #  parent context first and then child context. See:
              #  https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/control_flow_ops.py#L938
              if bwd_ctx:
                  for ctx in bwd_ctx:
                      _ = WhileContext(
                          context_def=ctx.to_proto(),
                          grad_state=ctx._grad_state,
                          import_scope=replica_prefix(i),
                      )
              if fwd_ctx:
                  for ctx in fwd_ctx:
                      _ = WhileContext(
                          context_def=ctx.to_proto(),
                          grad_state=ctx._grad_state,
                          import_scope=replica_prefix(i),
                      )
  
          # update saver
          master_replica = 0
          if graph_item.info.savers:
              item.info.update_savers(
                  [
                      Saver.from_proto(
                          proto,
                          import_scope=replica_prefix(master_replica),
                      ).to_proto()
                      for proto in graph_item.info.savers
                  ],
                  replace=False,
              )
  
          # update gradient info
          for i in range(self._num_local_replicas):
              for (
                  g_name,
                  t_name,
              ) in graph_item.grad_target_name_pairs.items():
                  if isinstance(g_name, tuple):
                      new_g_name = (
                          ops.prepend_name_scope(
                              g_name[0], replica_prefix(i)
                          ),
                          ops.prepend_name_scope(
                              g_name[1], replica_prefix(i)
                          ),
                          ops.prepend_name_scope(
                              g_name[2], replica_prefix(i)
                          ),
                      )
                  else:
                      new_g_name = ops.prepend_name_scope(
                          g_name, replica_prefix(i)
                      )
                  new_t_name = ops.prepend_name_scope(
                      t_name, replica_prefix(i)
                  )
                  item.extend_gradient_info_by_names(
                      grads=[new_g_name], targets=[new_t_name]
                  )
              item.info.update_variables(
                  [
                      _from_proto_fn(
                          proto, import_scope=replica_prefix(i)
                      ).to_proto()
                      for proto in graph_item.info.variables
                  ],
                  replace=False,
              )
              item.info.update_table_initializers(
                  [
                      ops.prepend_name_scope(tb_init, replica_prefix(i))
                      for tb_init in graph_item.info.table_initializers
                  ],
                  replace=False,
              )
      return item
  
file_path: autodist/autodist/kernel/replicator.py
imports: from tensorflow.python import ops, import_graph_def
  from tensorflow.python.framework import device_spec, kernels
  from tensorflow.python.framework.device_spec import DeviceSpecV2
  from tensorflow.python.ops.control_flow_ops import WhileContext
  from tensorflow.python.ops.control_flow_util import GetWhileContext
  from tensorflow.python.ops.resource_variable_ops import _from_proto_fn
  from autodist.graph_item import GraphItem
  from autodist.kernel.common.utils import replica_prefix
  from autodist.kernel.kernel import Kernel
  from autodist.checkpoint.saver import Saver
  from autodist.utils import logging
label_desc: Use image node_id 3 for calling the Replicator obj's underlying member method code with example usage: obj.replicate(graph_item) and returns: item