Copy an op directly to a given graph. Generally `op`'s inputs should already have been copied. If this is not the case, for example with v1 while_loops, then `_copy_non_source` inserts placeholders for the unavailable Tensors and returns a list of required mutations. Args: op: The op
(op, graph, op_map, base_graph)
| 60 | |
| 61 | |
| 62 | def _copy_non_source(op, graph, op_map, base_graph): |
| 63 | """Copy an op directly to a given graph. |
| 64 | |
| 65 | Generally `op`'s inputs should already have been copied. If this is not the |
| 66 | case, for example with v1 while_loops, then `_copy_non_source` inserts |
| 67 | placeholders for the unavailable Tensors and returns a list of required |
| 68 | mutations. |
| 69 | |
| 70 | Args: |
| 71 | op: The op to be copied. |
| 72 | graph: The destination graph. |
| 73 | op_map: A dict mapping ops and tensors in the old graph to the new one. |
| 74 | base_graph: The graph we're copying from, for any necessary functions. |
| 75 | Returns: |
| 76 | A tuple of (required_inputs, required_control_inputs): |
| 77 | required_inputs: |
| 78 | A list of `_InputMutation` tuples containing inputs to `copied_op` which |
| 79 | must be updated once `old_graph_tensor` has been copied. |
| 80 | required_control_inputs: |
| 81 | A list of `_ControlMutation` tuples containing control inputs to |
| 82 | `copied_op` which must be added once `old_graph_op` has been copied. |
| 83 | """ |
| 84 | input_mutations = [] |
| 85 | control_mutations = [] |
| 86 | copied_inputs = [] |
| 87 | for input_index, original_input in enumerate(op.inputs): |
| 88 | copied_input = op_map.get(original_input, None) |
| 89 | if copied_input is None: |
| 90 | # An input for this op is missing due to a loop in the graph. We'll insert |
| 91 | # a placeholder for now and return information about the required post-hoc |
| 92 | # mutation. |
| 93 | copied_input = array_ops.placeholder( |
| 94 | name="unused_control_flow_input", |
| 95 | shape=original_input.shape, |
| 96 | dtype=original_input.dtype) |
| 97 | input_mutations.append( |
| 98 | # `copied_op` is filled in below, after we've created it. |
| 99 | _InputMutation(copied_op=None, |
| 100 | input_index=input_index, |
| 101 | old_graph_tensor=original_input)) |
| 102 | copied_inputs.append(copied_input) |
| 103 | |
| 104 | copied_control_inputs = [] |
| 105 | for original_control_input in op.control_inputs: |
| 106 | copied_control_input = op_map.get(original_control_input, None) |
| 107 | if copied_control_input is None: |
| 108 | control_mutations.append( |
| 109 | _ControlMutation(copied_op=None, |
| 110 | old_graph_op=original_control_input)) |
| 111 | else: |
| 112 | copied_control_inputs.append(copied_control_input) |
| 113 | |
| 114 | # Don't copy over nodes with _tpu_replicate attribute. This attributed is used |
| 115 | # to signal that the op was built inside a tpu_replicate context; if we're |
| 116 | # lifting it to another graph we're similarly lifting it into another context. |
| 117 | with ops.control_dependencies(copied_control_inputs), ops.device(op.device): |
| 118 | # pylint: disable=protected-access |
| 119 | f = base_graph._functions.get(op.type, None) |
no test coverage detected