Returns the TPUReplicateContext and its associated graph.
()
| 164 | |
| 165 | |
| 166 | def _enclosing_tpu_context_and_graph(): |
| 167 | """Returns the TPUReplicateContext and its associated graph.""" |
| 168 | graph = ops.get_default_graph() |
| 169 | while graph is not None: |
| 170 | # pylint: disable=protected-access |
| 171 | context_ = graph._get_control_flow_context() |
| 172 | # pylint: enable=protected-access |
| 173 | while context_ is not None: |
| 174 | if isinstance(context_, TPUReplicateContext): |
| 175 | return context_, graph |
| 176 | context_ = context_.outer_context |
| 177 | graph = getattr(graph, "outer_graph", None) |
| 178 | raise ValueError("get_replicated_var_handle() called without " |
| 179 | "TPUReplicateContext. This shouldn't happen. Please file " |
| 180 | "a bug.") |
| 181 | |
| 182 | |
| 183 | class TPUReplicateContext(control_flow_ops.XLAControlFlowContext): |
no test coverage detected