(self, op_or_tensor)
| 1249 | self._conversion_map[old_output] = new_output |
| 1250 | |
| 1251 | def _convert_helper(self, op_or_tensor): |
| 1252 | stack = [op_or_tensor] |
| 1253 | while stack: |
| 1254 | y = stack[0] |
| 1255 | if y in self._conversion_map: |
| 1256 | assert isinstance(self._conversion_map[y], |
| 1257 | (WrappedTensor, ops.Operation)) |
| 1258 | stack.pop(0) |
| 1259 | continue |
| 1260 | if isinstance(y, ops.Operation): |
| 1261 | assert not y.outputs, ( |
| 1262 | "We only support converting Operation objects with no outputs. " |
| 1263 | "Got %s", y) |
| 1264 | y_op = y |
| 1265 | else: |
| 1266 | assert isinstance(y, ops.Tensor), y |
| 1267 | y_op = y.op |
| 1268 | |
| 1269 | is_while_loop = y_op.type == "Exit" |
| 1270 | if is_while_loop: |
| 1271 | while_op = WhileOp( |
| 1272 | y, pfor_ops=self._pfor_ops, pfor_config=self._pfor_config) |
| 1273 | is_inside_loop = while_op.is_inside_loop |
| 1274 | # If all nodes in the while_loop graph were created inside the pfor, we |
| 1275 | # treat the whole loop subgraph as a single op (y_op) and try to convert |
| 1276 | # it. For while_loops that are created completely or partially outside, |
| 1277 | # we treat them as external and should be able to simply return the Exit |
| 1278 | # node output as is without needing any conversion. Note that for |
| 1279 | # while_loops that are partially constructed inside, we assume they will |
| 1280 | # be loop invariant. If that is not the case, it will create runtime |
| 1281 | # errors since the converted graph would depend on the self._loop_var |
| 1282 | # placeholder. |
| 1283 | if is_inside_loop: |
| 1284 | y_op = while_op |
| 1285 | else: |
| 1286 | is_inside_loop = self.op_is_inside_loop(y_op) |
| 1287 | |
| 1288 | # If this op was not created inside the loop body, we will return as is. |
| 1289 | # 1. Convert inputs and control inputs. |
| 1290 | |
| 1291 | def _add_to_stack(x): |
| 1292 | if x not in self._conversion_map: |
| 1293 | stack.insert(0, x) |
| 1294 | return True |
| 1295 | else: |
| 1296 | return False |
| 1297 | |
| 1298 | if is_inside_loop: |
| 1299 | added_to_stack = False |
| 1300 | for inp in y_op.inputs: |
| 1301 | added_to_stack |= _add_to_stack(inp) |
| 1302 | for cinp in y_op.control_inputs: |
| 1303 | if cinp.outputs: |
| 1304 | for t in cinp.outputs: |
| 1305 | added_to_stack |= _add_to_stack(t) |
| 1306 | else: |
| 1307 | added_to_stack |= _add_to_stack(cinp) |
| 1308 | if added_to_stack: |
no test coverage detected