Initializer. Args: exit_node: A tensor output from the while_loop. pfor_ops: list of ops inside the current pfor loop. pfor_config: PForConfig object used while constructing loop body.
(self, exit_node, pfor_ops, pfor_config)
| 104 | """Object for storing state for converting the outputs of a while_loop.""" |
| 105 | |
| 106 | def __init__(self, exit_node, pfor_ops, pfor_config): |
| 107 | """Initializer. |
| 108 | |
| 109 | Args: |
| 110 | exit_node: A tensor output from the while_loop. |
| 111 | pfor_ops: list of ops inside the current pfor loop. |
| 112 | pfor_config: PForConfig object used while constructing loop body. |
| 113 | """ |
| 114 | self._pfor_config = pfor_config |
| 115 | self._pfor_ops = set(pfor_ops) |
| 116 | self._pfor_op_ids = set([x._id for x in pfor_ops]) |
| 117 | assert isinstance(exit_node, ops.Tensor) |
| 118 | self._while_context = exit_node.op._get_control_flow_context() |
| 119 | assert isinstance(self._while_context, control_flow_ops.WhileContext) |
| 120 | self._context_name = self._while_context.name |
| 121 | self._condition = self._while_context.pivot.op.inputs[0] |
| 122 | # Parts of an external while_loop could be created inside a pfor loop. |
| 123 | # However for the purpose here, we declare such loops to be external. Also |
| 124 | # note that we check if the condition was created inside or outside to |
| 125 | # determine if the while_loop was first created inside or outside. |
| 126 | # TODO(agarwal): check that the Enter and Exit of this loop are unstacked. |
| 127 | self._is_inside_loop = self.op_is_inside_loop(self._condition.op) |
| 128 | if self._is_inside_loop: |
| 129 | for e in self._while_context.loop_exits: |
| 130 | assert self.op_is_inside_loop(e.op) |
| 131 | |
| 132 | # Note the code below tries to reverse engineer an existing while_loop graph |
| 133 | # by assuming the following pattern of nodes. |
| 134 | # |
| 135 | # NextIteration <---- Body <--- Enter |
| 136 | # | ^ |
| 137 | # V ___| Y |
| 138 | # Enter -> Merge -> Switch___ |
| 139 | # ^ | N |
| 140 | # | V |
| 141 | # LoopCond Exit |
| 142 | |
| 143 | # Node that elements in the list below correspond one-to-one with each |
| 144 | # other. i.e. these lists are the same size, and the i_th entry corresponds |
| 145 | # to different Operations/Tensors of a single cycle as illustrated above. |
| 146 | # List of Switch ops (ops.Operation) that feed into an Exit Node. |
| 147 | self._exit_switches = [] |
| 148 | # List of inputs (ops.Tensor) to NextIteration. |
| 149 | self._body_outputs = [] |
| 150 | # List of list of control inputs of the NextIteration nodes. |
| 151 | self._next_iter_control_inputs = [] |
| 152 | # List of Merge ops (ops.Operation). |
| 153 | self._enter_merges = [] |
| 154 | # List of output (ops.Tensor) of Exit nodes. |
| 155 | self._outputs = [] |
| 156 | |
| 157 | # List of Enter Tensors. |
| 158 | # There are two types of Enter nodes: |
| 159 | # - The Enter nodes that are used in the `loop_vars` argument to |
| 160 | # `while_loop` (see |
| 161 | # https://www.tensorflow.org/api_docs/python/tf/while_loop). We collect |
| 162 | # these Enter nodes immediately below by tracing backwards from the Exit |
| 163 | # nodes via Exit <- Switch <- Merge <- Enter. You can see this chain in the |
nothing calls this directly
no test coverage detected