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","."]
code: def ExitGradientColocation(
      self, op: ops.Operation, gradient_uid: Text
  ):
      if op is not None:
          if (
              ops.get_default_graph()._control_flow_context is None
          ):  # pylint: disable=protected-access
              # Inside a TF2 tf.function or control flow graph and `op` was not
              # marked to be outside compiled.
              assert self._outside_compilation_v2_context is None
              return
          if self._outside_compilation_v2_context is not None:
              # Inside a TF2 tf.function or control flow graph and `op` was
              # marked to be outside compiled.
              self._outside_compilation_v2_context.Exit()
              self._outside_compilation_v2_context = None
              return
          if not self._gradient_colocation_stack:
              raise errors.InternalError(
                  op.node_def,
                  op,
                  (
                      "Badly nested gradient colocation: "
                      + f"empty stack when popping Op {op.name}"
                  ),
              )
          last_op = self._gradient_colocation_stack.pop()
          if op is last_op:
              if op is self._in_gradient_colocation:
                  self._in_gradient_colocation = None
                  self._ExitOutsideCompilationScope()
          else:
              raise errors.InternalError(
                  op.node_def,
                  op,
                  (
                      "Badly nested gradient colocation, "
                      + f"expected {last_op}, got {op.name}"
                  ),
              )
  
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 5 for calling the TPUReplicateContext obj's underlying member method code with example usage: obj.ExitGradientColocation(op, gradient_uid) without return types