Check if it is currently under `_TPUInferenceContext`.
()
| 1527 | |
| 1528 | |
| 1529 | def under_tpu_inference_context(): |
| 1530 | """Check if it is currently under `_TPUInferenceContext`.""" |
| 1531 | graph = ops.get_default_graph() |
| 1532 | while graph: |
| 1533 | context = graph._get_control_flow_context() # pylint: disable=protected-access |
| 1534 | while context: |
| 1535 | if isinstance(context, _TPUInferenceContext): |
| 1536 | return True |
| 1537 | context = context.outer_context |
| 1538 | if isinstance(graph, function._FuncGraph): # pylint: disable=protected-access |
| 1539 | graph = graph._outer_graph # pylint: disable=protected-access |
| 1540 | elif isinstance(graph, func_graph.FuncGraph): |
| 1541 | graph = graph.outer_graph |
| 1542 | else: |
| 1543 | return False |
| 1544 | |
| 1545 | |
| 1546 | class _TPUInferenceContext(control_flow_ops.XLAControlFlowContext): |
nothing calls this directly
no test coverage detected