class_docstr: A `ControlFlowContext` for nodes inside a TPU computation.
  
  The primary role of `TPUReplicateContext` is to mark operators inside a
  tpu.replicate() computation with the attribute "_tpu_replicate=XYZ", where XYZ
  is a unique name.
  
  We use a `ControlFlowContext` to perform the annotation since it integrates
  with Tensorflow constructs like ResourceVariables. For example, if a
  `ResourceVariable` is constructed inside a tpu.replicate() block, the
  `ResourceVariable` implementation can use
  `with ops.control_dependencies(None)` to build the variable's definition
  outside the replicated computation.
class_docstr_tok: ["A","`","ControlFlowContext","`","for","nodes","inside","a","TPU","computation",".","The","primary","role","of","`","TPUReplicateContext","`","is","to","mark","operators","inside","a","tpu.replicate","(",")","computation","with","the","attribute","``","_tpu_replicate=XYZ","''",",","where","XYZ","is","a","unique","name",".","We","use","a","`","ControlFlowContext","`","to","perform","the","annotation","since","it","integrates","with","Tensorflow","constructs","like","ResourceVariables",".","For","example",",","if","a","`","ResourceVariable","`","is","constructed","inside","a","tpu.replicate","(",")","block",",","the","`","ResourceVariable","`","implementation","can","use","`","with","ops.control_dependencies","(","None",")","`","to","build","the","variable","'s","definition","outside","the","replicated","computation","."]
docstr: Builds a new TPUReplicateContext.
  
  Args:
    name: a unique name for the context, used to populate the `_tpu_replicate`
      attribute.
    num_replicas: an integer that gives the number of replicas for the
      computation.
    pivot: a pivot node. Nodes in the TPUReplicateContext that do not have any
      inputs will have a control dependency on the pivot node. This ensures
      that nodes are correctly included in any enclosing control flow
      contexts.
docstr_tok: ["Builds","a","new","TPUReplicateContext",".","Args",":","name",":","a","unique","name","for","the","context",",","used","to","populate","the","`","_tpu_replicate","`","attribute",".","num_replicas",":","an","integer","that","gives","the","number","of","replicas","for","the","computation",".","pivot",":","a","pivot","node",".","Nodes","in","the","TPUReplicateContext","that","do","not","have","any","inputs","will","have","a","control","dependency","on","the","pivot","node",".","This","ensures","that","nodes","are","correctly","included","in","any","enclosing","control","flow","contexts","."]
code: def __init__(
      self, name: Text, num_replicas: int, pivot: ops.Operation
  ):
      """Builds a new TPUReplicateContext.
  
      Args:
        name: a unique name for the context, used to populate the `_tpu_replicate`
          attribute.
        num_replicas: an integer that gives the number of replicas for the
          computation.
        pivot: a pivot node. Nodes in the TPUReplicateContext that do not have any
          inputs will have a control dependency on the pivot node. This ensures
          that nodes are correctly included in any enclosing control flow
          contexts.
      """
      super(TPUReplicateContext, self).__init__()
      self._num_replicas = num_replicas
      self._outer_device_function_stack = None
      self._oc_dev_fn_stack = None
      self._outside_compilation_cluster = None
      self._is_map_outside_compilation = False
      self._outside_compilation_v2_context = None
      self._outside_compilation_counter = 0
      self._in_gradient_colocation = None
      self._gradient_colocation_stack = []
      self._host_compute_core = []
      self._name = name
      self._tpu_replicate_attr = attr_value_pb2.AttrValue(
          s=compat.as_bytes(self._name)
      )
      self._unsupported_ops = []
      self._pivot = pivot
      self._replicated_vars = {}
  
imports: from typing import Any, Callable, List, Optional, Text, Tuple, Union
  from absl import logging
  from tensorflow.core.framework import attr_value_pb2
  from tensorflow.python.distribute import device_util
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.framework import device
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import func_graph
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.tpu import device_assignment
  from tensorflow.python.tpu.ops import tpu_ops
  from tensorflow.python.types import core
  from tensorflow.python.util import compat
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 1 to create a new TPUReplicateContext object from inherited base classes: control_flow_ops with example: obj = TPUReplicateContext(name, num_replicas, pivot)