| 1325 | |
| 1326 | |
| 1327 | def control_op_remap(op, prefix, blob_remap): |
| 1328 | net_arg_names = [] |
| 1329 | if op.type == "If" or op.type == "AsyncIf": |
| 1330 | net_arg_names = ['then_net', 'else_net'] |
| 1331 | else: |
| 1332 | net_arg_names = ['loop_net', 'cond_net'] |
| 1333 | for argument in op.arg: |
| 1334 | if argument.name in net_arg_names: |
| 1335 | assert argument.n, \ |
| 1336 | "Expected non empty net in " + op.type + "'s " + argument.name + " argument" |
| 1337 | subnet = Net(argument.n) |
| 1338 | remapped_subnet = subnet.Clone( |
| 1339 | name=(subnet._net.name if subnet._net.name else '') + '_remapped', |
| 1340 | blob_remap=blob_remap) |
| 1341 | argument.n.CopyFrom(remapped_subnet.Proto()) |
| 1342 | |
| 1343 | |
| 1344 | DEFAULT_REMAP_FUNCS = { |