Initialize the pending count for ops between two lists of Operations. 'pending_count[op]' indicates the number of backprop inputs to this operation. Args: to_ops: list of Operations. from_ops: list of Operations. colocate_gradients_with_ops: Python bool. See docstring of gradien
(to_ops, from_ops, colocate_gradients_with_ops, func_graphs,
xs_set)
| 69 | |
| 70 | |
| 71 | def _PendingCount(to_ops, from_ops, colocate_gradients_with_ops, func_graphs, |
| 72 | xs_set): |
| 73 | """Initialize the pending count for ops between two lists of Operations. |
| 74 | |
| 75 | 'pending_count[op]' indicates the number of backprop inputs |
| 76 | to this operation. |
| 77 | |
| 78 | Args: |
| 79 | to_ops: list of Operations. |
| 80 | from_ops: list of Operations. |
| 81 | colocate_gradients_with_ops: Python bool. See docstring of gradients(). |
| 82 | func_graphs: list of FuncGraphs. This method will traverse through |
| 83 | these functions if they capture from_ops or any reachable ops. This is |
| 84 | useful if to_ops occur in a function and from_ops are in an outer function |
| 85 | or graph. |
| 86 | xs_set: ObjectIdentitySet of Tensors. |
| 87 | |
| 88 | Returns: |
| 89 | A tuple containing: (1) the subset of to_ops reachable from from_ops by a |
| 90 | path of zero or more backpropagatable tensors, (2) a mapping from operation |
| 91 | to the number of backprop inputs to that op, and (3) a ControlFlowState |
| 92 | object which is not None if the ops between from_ops and to_ops contain |
| 93 | control flow loops. |
| 94 | """ |
| 95 | # Mark reachable ops from from_ops. |
| 96 | reached_ops = set() |
| 97 | _MarkReachedOps(from_ops, reached_ops, func_graphs) |
| 98 | # X in reached_ops iff X is reachable from from_ops by a path of zero or more |
| 99 | # backpropagatable tensors. |
| 100 | |
| 101 | reachable_to_ops = set(op for op in to_ops if op in reached_ops) |
| 102 | |
| 103 | # Mark between ops. |
| 104 | between_ops = set() |
| 105 | between_op_list = [] |
| 106 | queue = collections.deque() |
| 107 | queue.extend(to_ops) |
| 108 | while queue: |
| 109 | op = queue.popleft() |
| 110 | # We are interested in this op. |
| 111 | if op in reached_ops: |
| 112 | between_ops.add(op) |
| 113 | between_op_list.append(op) |
| 114 | # Clear the boolean so we won't add the inputs again. |
| 115 | reached_ops.remove(op) |
| 116 | for inp in _NonEagerInputs(op, xs_set): |
| 117 | queue.append(inp.op) |
| 118 | # X in between_ops iff X is on a path of zero or more backpropagatable tensors |
| 119 | # between from_ops and to_ops |
| 120 | |
| 121 | # 'loop_state' is None if there are no while loops. |
| 122 | loop_state = control_flow_state.MaybeCreateControlFlowState( |
| 123 | between_op_list, between_ops, colocate_gradients_with_ops) |
| 124 | |
| 125 | # Initialize pending count for between ops. |
| 126 | pending_count = collections.defaultdict(int) |
| 127 | for op in between_op_list: |
| 128 | for x in _NonEagerInputs(op, xs_set): |
no test coverage detected