()
| 810 | # TODO(b/138808492): Remove code inside copybara |
| 811 | # to make TPU code and CPU code consistent. |
| 812 | def _enclosing_tpu_context(): |
| 813 | # pylint: disable=protected-access |
| 814 | context = ops.get_default_graph()._get_control_flow_context() |
| 815 | # pylint: enable=protected-access |
| 816 | while context is not None and not isinstance( |
| 817 | context, control_flow_ops.XLAControlFlowContext): |
| 818 | context = context.outer_context |
| 819 | return context |
| 820 | |
| 821 | |
| 822 | # copybara:strip_end |
no test coverage detected