| 381 | graph._device_function_stack = self._oc_dev_fn_stack # pylint: disable=protected-access |
| 382 | |
| 383 | def Enter(self): |
| 384 | if not self._outer_device_function_stack: |
| 385 | # Capture the device function stack at the time of first entry |
| 386 | # since that is the stack that will be used outside_compilation. |
| 387 | graph = ops.get_default_graph() |
| 388 | # pylint: disable=protected-access |
| 389 | self._outer_device_function_stack = graph._device_function_stack.copy() |
| 390 | # pylint: enable=protected-access |
| 391 | super(TPUReplicateContext, self).Enter() |
| 392 | |
| 393 | def HostComputeCore(self): |
| 394 | return self._host_compute_core |