Function outside_compilation

tensorflow/python/tpu/tpu.py:536–575

Builds part of a computation outside any current TPU replicate scope. 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:

(computation, *args, **kwargs)

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

  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 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


@tf_export(v1=["tpu.replicate"])