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 EnterGradientColocation(
      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
              # If we are in TF 2 functions (control flow V2 functions, or
              # tf.function()), we need to attach _xla_outside_compilation attribute
              # directly because we are not in TPUReplicateContext.
              try:
                  outside_attr = op.get_attr(
                      _OUTSIDE_COMPILATION_ATTR
                  ).decode("ascii")
              except ValueError:
                  # The attr was not present: do nothing.
                  return
              parts = outside_attr.split(".")
              cluster = parts[0] + "." + gradient_uid
              self._outside_compilation_v2_context = (
                  OutsideCompilationV2Context(cluster)
              )
              self._outside_compilation_v2_context.Enter()
              return
          self._gradient_colocation_stack.append(op)
          if not self._outside_compilation_cluster:
              try:
                  outside_attr = op.get_attr(
                      _OUTSIDE_COMPILATION_ATTR
                  ).decode("ascii")
                  if self._in_gradient_colocation:
                      raise NotImplementedError(
                          "Cannot nest gradient colocation operations outside compilation"
                      )
                  if gradient_uid == "__unsupported__":
                      raise NotImplementedError(
                          "No gradient_uid calling gradient within outside_compilation"
                      )
                  # When we take the gradient of an op X in an outside_compilation
                  # cluster C in a forward computation we would like to put the ops
                  # corresponding to the gradient of X into a new outside_compilation
                  # cluster C'. However, if we take the gradient of X twice, the second
                  # one should get yet another new outside_compilation cluster C''.
                  #
                  # The mechanism we adopt is to use a 'root_cluster' which is the
                  # cluster that X was in before we took gradients, and a 'gradient_uid'
                  # which is different for every invocation of gradients, and put the
                  # gradient of X in cluster 'root_cluster.gradient_uid'.
                  #
                  # When taking a gradient of a gradient, some ops will be colocated
                  # with Op in the forward pass (e.g., cluster root_cluster) and some in
                  # the backward pass (e.g., cluster root_cluster.initial_gradient_uid).
                  # We need all of the grad-of-grad ops to be in the same cluster to
                  # avoid cyclic dependencies between clusters. We adopt a heuristic
                  # that puts any op clustered with root_cluster.<xxx> in
                  # root_cluster.gradient_uid, even if xxx was initial_gradient_uid.
                  self._in_gradient_colocation = op
                  parts = outside_attr.split(".")
                  cluster = parts[0] + "." + gradient_uid
                  self._EnterOutsideCompilationScope(cluster=cluster)
              except ValueError:
                  # The attr was not present: do nothing.
                  pass
  
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 4 for calling the TPUReplicateContext obj's underlying member method code with example usage: obj.EnterGradientColocation(op, gradient_uid) without return types