# 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 distribute_lib 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" _MAP_OUTSIDE_COMPILATION_ATTR = "_xla_map_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 distribute_lib.has_strategy(): return None strategy = distribute_lib.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._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 = {} 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. 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, is_map_outside_compilation=False ): 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 if is_map_outside_compilation: self._is_map_outside_compilation = True 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 self._is_map_outside_compilation = False 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._is_map_outside_compilation: op._set_attr( _MAP_OUTSIDE_COMPILATION_ATTR, attr_value_pb2.AttrValue(b=True), ) 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, is_map_outside_compilation=False): control_flow_ops.ControlFlowContext.__init__(self) self._name = name self._is_map_outside_compilation = is_map_outside_compilation def AddOp(self, op: ops.Operation) -> None: if self._outer_context: self._outer_context.AddOp(op) self._set_outside_compilation_attributes(op) def AddInnerOp(self, op: ops.Operation) -> None: if self._outer_context: self._outer_context.AddInnerOp(op) self._set_outside_compilation_attributes(op) def to_control_flow_context_def(self, context_def, export_scope=None): raise NotImplementedError def _set_outside_compilation_attributes(self, op: ops.Operation) -> None: # pylint: disable=protected-access op._set_attr( _OUTSIDE_COMPILATION_ATTR, attr_value_pb2.AttrValue(s=compat.as_bytes(self._name)), ) if self._is_map_outside_compilation: op._set_attr( _MAP_OUTSIDE_COMPILATION_ATTR, attr_value_pb2.AttrValue(b=True) ) # pylint: enable=protected-access def outside_compilation_impl( is_map, computation: Callable[..., Any], *args, **kwargs ) -> Any: """Tags ops in `computation` with outside compilation attributes for ordinary `outside_compilation` or `map_outside_compilation`.""" 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, is_map_outside_compilation=is_map ) 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(is_map_outside_compilation=is_map) # 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.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, 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 a later passes ops with outside compilation attributes are moved to a host-side graph. Inputs to this extract host-side graph are 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. """ return outside_compilation_impl(False, computation, *args, **kwargs) def experimental_map_outside_compilation( computation: Callable[..., Any], *args, **kwargs ) -> Any: """Maps `computation` onto shards and puts it outside any current TPU replicate scope. `experimental_map_outside_compilation(f, x)` maps `f` onto the shards of `x`, where `x` is split-sharded. Each invocation of `f` on a split occurs on the CPU that's associated with the TPU that owns the split. Example usage: ```python def normalize_each_split(split): return split - tf.math.reduce_mean(split) def tpu_computation(x): x_split = strategy.experimental_split_to_logical_devices( x, [num_cores_per_replica, 1]) y = experimental_map_outside_compilation( normalize_each_split, x_split) y_split = strategy.experimental_split_to_logical_devices( x, [num_cores_per_replica, 1]) return y_split ``` `experimental_map_outside_compilation` should be called inside TPUReplicateContext. That is, `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()`. It is invalid to invoke outside of TPUReplicateContext. `experimental_map_outside_compilation` should input and output tensors that are located on the TPU. Internally, `experimental_map_outside_compilation()` adds outside compilation attributes to all ops in `computation` and moves outside-compiled ops to a host-side graph. This is similar to `tf.tpu.outside_compilation()`. Send/recv ops from/to the TPU send each split directly to the TPU's host. 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. """ return outside_compilation_impl(True, computation, *args, **kwargs)