Returns the XLAControlFlowContext, which exists inside a tpu.rewrite().
()
| 1089 | |
| 1090 | |
| 1091 | def _enclosing_tpu_context(): |
| 1092 | """Returns the XLAControlFlowContext, which exists inside a tpu.rewrite().""" |
| 1093 | graph = ops.get_default_graph() |
| 1094 | while graph is not None: |
| 1095 | # pylint: disable=protected-access |
| 1096 | context_ = graph._get_control_flow_context() |
| 1097 | # pylint: enable=protected-access |
| 1098 | while context_ is not None: |
| 1099 | if isinstance(context_, control_flow_ops.XLAControlFlowContext): |
| 1100 | return context_ |
| 1101 | context_ = context_.outer_context |
| 1102 | # This may be a FuncGraph due to defuns or v2 control flow. We need to |
| 1103 | # find the original graph with the XLAControlFlowContext. |
| 1104 | graph = getattr(graph, "outer_graph", None) |
| 1105 | return None |
| 1106 | |
| 1107 | |
| 1108 | def is_distributed_variable(v): |
no test coverage detected