tpu_replication.py

# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file8 except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ======================================

"""OutsideCompilation, TPUReplicateContext, and supporting functions."""

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 distribution_strategy_context
from tensorflow.python.framework import device as pydev
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 as device_assignment_lib
from tensorflow.python.tpu.ops import tpu_ops
from tensorflow.python.types import core as core_types
from tensorflow.python.util import compat
from tensorflow.python.util.tf_export import tf_export

_MAX_WARNING_LINES = 5
_TPU_REPLICATE_ATTR = "_tpu_replicate"
_OUTSIDE_COMPILATION_ATTR = "_xla_outside_compilation"

# Operations that indicate some error in the users graph, e.g. a placeholder
# that's introduced outside of the infeed.
_DENYLISTED_OPS = frozenset([
    "Placeholder",
])


# XLA doesn't currently support reading of intermediate tensors, thus some ops
# are not supported.
_UNSUPPORTED_OPS = frozenset([
    "AudioSummary",
    "AudioSummaryV2",
    "HistogramSummary",
    "ImageSummary",
    "MergeSummary",
    "Print",
    "ScalarSummary",
    "TensorSummary",
    "TensorSummaryV2",
])


def is_tpu_strategy(strategy: Any) -> bool:
  is_tpu_strat = lambda k: k.__name__.startswith("TPUStrategy")
  clz = strategy.__class__
  return is_tpu_strat(clz) or any(map(is_tpu_strat, clz.__bases__))


def _enclosing_tpu_device_assignment(
) -> Optional[device_assignment_lib.DeviceAssignment]:
  if not distribution_strategy_context.has_strategy():
    return None
  strategy = distribution_strategy_context.get_strategy()
  if not is_tpu_strategy(strategy):
    return None
  return strategy.extended._device_assignment  # pylint: disable=protected-access


class TPUReplicateContext(control_flow_ops.XLAControlFlowContext):
  """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.
  """

  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._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 = {}

  def get_replicated_var_handle(self,
                                name: Text,
                                handle_id: Text,
                                vars_: Union[List[core_types.Tensor],
                                             List[variables.Variable]],
                                is_mirrored: bool = False,
                                is_packed: bool = False) -> core_types.Tensor:
    """Returns a variable handle for replicated TPU variable 'var'.

    This is a method used by an experimental replicated variable implementation
    and is not intended as a public API.

    Args:
      name: The common name of the variable.
      handle_id: Unique ID of the variable handle, used as the cache key.
      vars_: The replicated TPU variables or handles.
      is_mirrored: Whether the variables are mirrored, which guarantees the
        values in each replica are always the same.
      is_packed: Whether the replicated variables are packed into one variable.

    Returns:
      The handle of the TPU replicated input node.
    """
    device_assignment = _enclosing_tpu_device_assignment()
    # We don't need to put device assignment as part of the replicated_vars key
    # because each TPUReplicateContext will only have one device assignment.
    handle = self._replicated_vars.get(handle_id)
    if handle is not None:
      return handle

    if device_assignment is not None and not is_packed:
      # Find a variable copy for each replica in the device assignment.
      # Note that the order of devices for replicas for the variable and the
      # device assignment might not match.
      job_name = pydev.DeviceSpec.from_string(vars_[0].device).job
      devices_to_vars = {device_util.canonicalize(v.device): v for v in vars_}
      replicated_vars = []
      for replica_id in range(device_assignment.num_replicas):
        for logical_core in range(device_assignment.num_cores_per_replica):
          device = device_util.canonicalize(
              device_assignment.tpu_device(
                  replica=replica_id, logical_core=logical_core, job=job_name))
          if device in devices_to_vars:
            replicated_vars.append(devices_to_vars[device])
            break
        else:
          raise ValueError(
              "Failed to find a variable on any device in replica {} for "
              "current device assignment".format(replica_id)
          )
    else:
      replicated_vars = vars_

    # Builds a TPUReplicatedInput node for the variable, if one does not already
    # exist. The TPUReplicatedInput node must belong to the enclosing
    # control-flow scope of the TPUReplicateContext.
    # TODO(phawkins): consider changing the contract of the TPU encapsulation
    # so the TPUReplicatedInput nodes go inside the TPUReplicateContext scope
    # instead.

    _, graph = _enclosing_tpu_context_and_graph()
    with graph.as_default():
      # If replicated_vars are variables, get the handles. Note that this can be
      # done inside TPUReplicateContext because replicated_vars.handle may
      # create new ops.
      if isinstance(replicated_vars[0], variables.Variable):
        replicated_vars = [v.handle for v in replicated_vars]
      # pylint: disable=protected-access
      saved_context = graph._get_control_flow_context()
      graph._set_control_flow_context(self.outer_context)
      handle = tpu_ops.tpu_replicated_input(
          replicated_vars,
          name=name + "/handle",
          is_mirrored_variable=is_mirrored,
          is_packed=is_packed)
      graph._set_control_flow_context(saved_context)
      # pylint: enable=protected-access
    self._replicated_vars[handle_id] = handle
    return handle

  def report_unsupported_operations(self) -> None:
    if self._unsupported_ops:
      op_str = "\n".join(
          "  %s (%s)" % (op.type, op.name) for op in
          self._unsupported_ops[:_MAX_WARNING_LINES])
      logging.warning("%d unsupported operations found: \n%s",
                      len(self._unsupported_ops), op_str)
      if len(self._unsupported_ops
            ) > _MAX_WARNING_LINES:
        logging.warning("... and %d more",
                        (len(self._unsupported_ops) - _MAX_WARNING_LINES))

  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

  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}")
        )

  def _EnterOutsideCompilationScope(self, cluster: Optional[Text] = None):

    class FakeOp(object):
      """A helper class to determine the current device.

      Supports only the type and device set/get methods needed to run the
      graph's _apply_device_function method.
      """

      def __init__(self):
        self._device = ""

      @property
      def type(self):
        return "FakeOp"

      @property
      def device(self):
        return self._device

      def _set_device(self, device):
        if isinstance(device, pydev.DeviceSpec):
          self._device = device.to_string()
        else:
          self._device = device

      def _set_device_from_string(self, device_str):
        self._device = device_str

    if self._outside_compilation_cluster:
      raise NotImplementedError("Cannot nest outside_compilation clusters")
    if cluster:
      self._outside_compilation_cluster = cluster
    else:
      self._outside_compilation_cluster = str(self._outside_compilation_counter)
      self._outside_compilation_counter += 1
    graph = ops.get_default_graph()
    fake_op = FakeOp()
    graph._apply_device_functions(fake_op)  # pylint: disable=protected-access
    device = pydev.DeviceSpec.from_string(fake_op.device)
    if (device.device_type == "TPU_REPLICATED_CORE" and
        device.device_index is not None):
      self._host_compute_core.append(self._outside_compilation_cluster + ":" +
                                     str(device.device_index))
    self._oc_dev_fn_stack = graph._device_function_stack  # pylint: disable=protected-access
    graph._device_function_stack = self._outer_device_function_stack  # pylint: disable=protected-access

  def _ExitOutsideCompilationScope(self):
    if not self._outside_compilation_cluster:
      raise ValueError(
          "Attempted to exit outside_compilation scope when not in scope")
    self._outside_compilation_cluster = None
    graph = ops.get_default_graph()
    graph._device_function_stack = self._oc_dev_fn_stack  # pylint: disable=protected-access

  def Enter(self) -> None:
    if not self._outer_device_function_stack:
      # Capture the device function stack at the time of first entry
      # since that is the stack that will be used outside_compilation.
      graph = ops.get_default_graph()
      # pylint: disable=protected-access
      self._outer_device_function_stack = graph._device_function_stack.copy()
      # pylint: enable=protected-access
    super(TPUReplicateContext, self).Enter()

  def HostComputeCore(self) -> List[Text]:
    return self._host_compute_core

  def _RemoveExternalControlEdges(
      self,
      op: ops.Operation) -> Tuple[List[ops.Operation], List[ops.Operation]]:
    """Remove any external control dependency on this op."""
    internal_control_inputs = []
    external_control_inputs = []
    for x in op.control_inputs:
      # pylint: disable=protected-access
      is_internal_op = False
      ctxt = x._get_control_flow_context()
      while ctxt is not None:
        if ctxt == self:
          is_internal_op = True
          break
        ctxt = ctxt._outer_context
      if is_internal_op:
        internal_control_inputs.append(x)
      else:
        external_control_inputs.append(x)
      # pylint: enable=protected-access
    # pylint: disable=protected-access
    op._remove_all_control_inputs()
    op._add_control_inputs(internal_control_inputs)
    # pylint: enable=protected-access
    return internal_control_inputs, external_control_inputs

  def AddOp(self, op: ops.Operation) -> None:
    # pylint: disable=protected-access
    if op.type in _DENYLISTED_OPS:
      logging.error(
          "Operation of type %s (%s) is not supported on the TPU. "
          "Execution will fail if this op is used in the graph. ", op.type,
          op.name)

    if op.type in _UNSUPPORTED_OPS:
      self._unsupported_ops.append(op)

    if any(x.dtype._is_ref_dtype for x in op.inputs):
      raise NotImplementedError(
          f"Non-resource Variables are not supported inside TPU computations "
          f"(operator name: {op.name})")

    # TensorFlowOpLayer may clone nodes that are in tpu.rewrite()s. It'll add
    # the "_cloned" attribute and we should continue in that case.
    if (_TPU_REPLICATE_ATTR in op.node_def.attr and
        "_cloned" not in op.node_def.attr):
      raise ValueError(f"TPU computations cannot be nested on op ({op})")
    op._set_attr(_TPU_REPLICATE_ATTR, self._tpu_replicate_attr)
    if self._outside_compilation_cluster:
      op._set_attr(
          _OUTSIDE_COMPILATION_ATTR,
          attr_value_pb2.AttrValue(
              s=compat.as_bytes(self._outside_compilation_cluster)))
    if self._num_replicas > 1 or not self._outside_compilation_cluster:
      # Prevent feeding or fetching anything that is being compiled,
      # and any replicated outside_compilation Op.
      op.graph.prevent_feeding(op)
      op.graph.prevent_fetching(op)

    # Remove any control edges from outer control flow contexts. These may cause
    # mismatched frame errors.
    (internal_control_inputs,
     external_control_inputs) = self._RemoveExternalControlEdges(op)

    if not op.inputs:
      # Add a control edge from the control pivot to this op.
      if not internal_control_inputs:
        # pylint: disable=protected-access
        op._add_control_input(self.GetControlPivot())
        # pylint: enable=protected-access
    else:
      for index in range(len(op.inputs)):
        x = op.inputs[index]
        real_x = self.AddValue(x)
        if real_x is not x:
          op._update_input(index, real_x)  # pylint: disable=protected-access

    if external_control_inputs:
      # Use an identity to pull control inputs as data inputs. Note that we
      # ignore ops which don't have outputs. TODO(phawkins): fix that.
      with ops.control_dependencies(None):
        self.Enter()
        external_control_inputs = [
            array_ops.identity(x.outputs[0]).op
            for x in external_control_inputs
            if x.outputs
        ]
        self.Exit()
      # pylint: disable=protected-access
      op._add_control_inputs(external_control_inputs)
      # pylint: enable=protected-access

    # Mark op's outputs as seen by this context and any outer contexts.
    output_names = [x.name for x in op.outputs]
    context = self
    while context is not None:
      # pylint: disable=protected-access
      context._values.update(output_names)
      context = context._outer_context
      # pylint: enable=protected-access

    if self._outer_context:
      self._outer_context.AddInnerOp(op)

  def AddValue(self, val: core_types.Tensor) -> core_types.Tensor:
    """Add `val` to the current context and its outer context recursively."""
    if not self._outer_context:
      return val

    if val.name in self._values:
      # Use the real value if it comes from outer context.
      result = self._external_values.get(val.name)
      return val if result is None else result

    result = val
    self._values.add(val.name)
    if self._outer_context:
      result = self._outer_context.AddValue(val)
      self._values.add(result.name)

    self._external_values[val.name] = result

    return result

  def AddInnerOp(self, op: ops.Operation):
    self.AddOp(op)
    if self._outer_context:
      self._outer_context.AddInnerOp(op)

  @property
  def grad_state(self):
    # Define the gradient loop state associated with the TPUReplicateContext to
    # be None as the TPUReplicateContext does not get nested nor does the
    # grad_state outside the TPUReplicateContext affect the graph inside so the
    # grad_state should be as if this is the top-level gradient state.
    return None

  @property
  def back_prop(self):
    """Forwards to the enclosing while context, if any."""
    if self.GetWhileContext():
      return self.GetWhileContext().back_prop
    return False

  def GetControlPivot(self) -> ops.Operation:
    return self._pivot

  def RequiresUniqueFunctionRetracing(self):
    # More context: b/158152827. TPU stack uses the TPUReplicateContext to
    # create replicated variable handles and cluster TPU computations, thus we
    # always retrace a tf.function when the wrapped TPUReplicateContext changes.
    return True


def _enclosing_tpu_context_and_graph() -> Tuple[Any, Any]:
  """Returns the TPUReplicateContext and its associated graph."""
  graph = ops.get_default_graph()
  while graph is not None:
    # pylint: disable=protected-access
    context_ = graph._get_control_flow_context()
    # pylint: enable=protected-access
    while context_ is not None:
      if isinstance(context_, TPUReplicateContext):
        return context_, graph
      context_ = context_.outer_context
    graph = getattr(graph, "outer_graph", None)
  raise ValueError("get_replicated_var_handle() called without "
                   "TPUReplicateContext. This shouldn't happen. Please file "
                   "a bug.")


class OutsideCompilationV2Context(control_flow_ops.ControlFlowContext):
  """The context for outside compilation in Tensorflow 2.0.

  Every op added in this context will be assigned an _xla_outside_compilation
  attribute.
  """

  def __init__(self, name: Text):
    control_flow_ops.ControlFlowContext.__init__(self)
    self._name = name

  def AddOp(self, op: ops.Operation) -> None:
    if self._outer_context:
      self._outer_context.AddOp(op)
    # pylint: disable=protected-access
    op._set_attr("_xla_outside_compilation",
                 attr_value_pb2.AttrValue(s=compat.as_bytes(self._name)))
    # pylint: enable=protected-access

  def AddInnerOp(self, op: ops.Operation) -> None:
    if self._outer_context:
      self._outer_context.AddInnerOp(op)
    # pylint: disable=protected-access
    op._set_attr("_xla_outside_compilation",
                 attr_value_pb2.AttrValue(s=compat.as_bytes(self._name)))
    # pylint: enable=protected-access

  def to_control_flow_context_def(self, context_def, export_scope=None):
    raise NotImplementedError


@tf_export(v1=["tpu.outside_compilation"])
def outside_compilation(computation: Callable[..., Any], *args,
                        **kwargs) -> Any:
  """Builds part of a computation outside any current TPU replicate scope.

  `tf.tpu.outside_compilation()` is used to run ops in `computation` on CPU
  instead of running on TPU. For example, users can run ops that are not
  supported on TPU's (e.g. tf.summary.write()) by explicitly placing those
  ops on CPU's. Below usage of outside compilation will place ops in
  `computation_with_string_ops` on CPU.

  Example usage:

  ```python
  def computation_with_string_ops(x):
    # strings types are not supported on TPU's and below ops must
    # run on CPU instead.
    output = tf.strings.format('1{}', x)
    return tf.strings.to_number(output)

  def tpu_computation():
    # Expected output is 11.
    output = tf.tpu.outside_compilation(computation_with_string_ops, 1)
  ```

  Outside compilation should be called inside TPUReplicateContext. That is,
  `tf.tpu.outside_compilation()` should be called inside a function that is
  passed to `tpu.split_compile_and_replicate()` -- this is implied when
  outside compilation is invoked inside a function passed to TPUStrategy
  `run()`. If invoked outside of TPUReplicateContext,
  then this simply returns the result of `computation`, and therefore,
  would be a no-op. Note that outside compilation is different from
  `tf.distribute.experimental.TPUStrategy.merge_call()` as logic in
  outside compilation is replicated and executed separately for each
  replica. On the other hand, `merge_call()` requires a `merge_fn`
  to aggregate the inputs from different replicas and is executed only
  once.

  For variables placed in TPU device, which includes variables created inside
  TPUStrategy scope, outside compilation logic must not include variable
  read/write. For variables placed on host, which is the case when variables
  created via TPUEstimator, variable read/write is only allowed if the variable
  is not accessed by any other ops in the TPU computation. Variable read/write
  from outside compilation cluster is not visible from TPU computation and
  vice versa. Therefore, if outside compilation logic contains such host
  variables read/write ops and if the variables are accessed by TPU
  computation as well, then this may lead to deadlock.

  Internally, `tf.tpu.outside_compilation()` adds outside compilation
  attributes to all ops in `computation`. During later graph pass, these
  ops with outside compilation attribute is extracted out and replicated
  into a host-side graph. Inputs to this extract host-side graph is sent
  from TPU computation graph to host graph via a pair of XlaSendToHost and
  XlaRecvFromHost ops. Note that using `tf.tpu.outside_compilation()`
  may result in tensor transfer between TPU and CPU, leading to non-trivial
  performance impact.

  Args:
    computation: A Python function that builds the computation to place on the
      host.
    *args: the positional arguments for the computation.
    **kwargs: the keyword arguments for the computation.

  Returns:
    The Tensors returned by computation.
  """
  args = [] if args is None else args
  graph = ops.get_default_graph()

  # 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.
  if isinstance(graph, func_graph.FuncGraph):
    try:
      tpu_context, _ = _enclosing_tpu_context_and_graph()
    except ValueError:
      logging.warning(
          "Outside compilation attempted outside TPUReplicateContext "
          "scope. As no enclosing TPUReplicateContext can be found, "
          "returning the result of `computation` as is.")
      return computation(*args, **kwargs)

    # pylint: disable=protected-access
    outside_compilation_name = str(tpu_context._outside_compilation_counter)
    tpu_context._outside_compilation_counter = (
        tpu_context._outside_compilation_counter + 1)
    # pylint: enable=protected-access

    outside_compilation_context = OutsideCompilationV2Context(
        outside_compilation_name)
    outside_compilation_context.Enter()
    args = [] if args is None else args
    retval = computation(*args, **kwargs)
    outside_compilation_context.Exit()
    return retval

  # If we are in a TPUReplicateContext, signal that we are now
  # outside_compilation
  initial_context = graph._get_control_flow_context()  # pylint: disable=protected-access
  context = initial_context
  while context:
    if isinstance(context, TPUReplicateContext):
      context._EnterOutsideCompilationScope()  # pylint: disable=protected-access
    context = context.outer_context

  retval = computation(*args, **kwargs)

  # If we are in a TPUReplicateContext, signal that we are no longer
  # outside_compilation
  final_context = graph._get_control_flow_context()  # pylint: disable=protected-access
  if initial_context is not final_context:
    raise NotImplementedError(
        "Control-flow context cannot be different at start and end of an "
        "outside_compilation scope")
  context = initial_context
  while context:
    if isinstance(context, TPUReplicateContext):
      context._ExitOutsideCompilationScope()  # pylint: disable=protected-access
    context = context.outer_context

  return retval