Core: Add the loop termination condition and body to the graph.
(self, pred, body, original_loop_vars, loop_vars,
shape_invariants)
| 2095 | raise TypeError("Type %s not supported" % type(x)) |
| 2096 | |
| 2097 | def _BuildLoop(self, pred, body, original_loop_vars, loop_vars, |
| 2098 | shape_invariants): |
| 2099 | """Core: Add the loop termination condition and body to the graph.""" |
| 2100 | flat_loop_vars = nest.flatten(original_loop_vars, expand_composites=True) |
| 2101 | |
| 2102 | # Let the context know the loop variables so the loop variables |
| 2103 | # would be added in the outer contexts properly. |
| 2104 | self._InitializeValues(loop_vars) |
| 2105 | real_vars = loop_vars |
| 2106 | if self._outer_context: |
| 2107 | real_vars = [self._outer_context.AddValue(x) for x in loop_vars] |
| 2108 | with ops.control_dependencies(None): |
| 2109 | enter_vars = [ |
| 2110 | _Enter( |
| 2111 | x, |
| 2112 | self._name, |
| 2113 | is_constant=False, |
| 2114 | parallel_iterations=self._parallel_iterations, |
| 2115 | use_input_shape=(shape_invariants is None)) for x in real_vars |
| 2116 | ] |
| 2117 | for x in enter_vars: |
| 2118 | x.graph.prevent_feeding(x) |
| 2119 | if self._outer_context: |
| 2120 | self._outer_context.AddInnerOp(x.op) |
| 2121 | |
| 2122 | # Finds the closest enclosing non-None control pivot. |
| 2123 | outer_context = self._outer_context |
| 2124 | control_pivot = None |
| 2125 | while outer_context is not None and control_pivot is None: |
| 2126 | control_pivot = outer_context.GetControlPivot() |
| 2127 | # pylint: disable=protected-access |
| 2128 | outer_context = outer_context._outer_context |
| 2129 | # pylint: enable=protected-access |
| 2130 | |
| 2131 | if control_pivot is not None: |
| 2132 | for var in enter_vars: |
| 2133 | if util.IsLoopConstantEnter(var.op.inputs[0].op): |
| 2134 | # pylint: disable=protected-access |
| 2135 | var.op._add_control_input(control_pivot.op) |
| 2136 | # pylint: enable=protected-access |
| 2137 | _SetShapeInvariants(real_vars, enter_vars, shape_invariants) |
| 2138 | |
| 2139 | # Fix the control inputs and control flow context of these enter ops. |
| 2140 | self._FixControlInputsAndContext(enter_vars) |
| 2141 | self._InitializeValues(enter_vars) |
| 2142 | self._loop_enters = enter_vars |
| 2143 | |
| 2144 | merge_vars = [merge([x, x])[0] for x in enter_vars] |
| 2145 | self._pivot_for_pred = merge_vars[0] |
| 2146 | |
| 2147 | # Build the graph for pred. |
| 2148 | merge_vars_with_tensor_arrays = ( |
| 2149 | _convert_flows_to_tensorarrays(flat_loop_vars, merge_vars)) |
| 2150 | packed_vars = nest.pack_sequence_as( |
| 2151 | structure=original_loop_vars, |
| 2152 | flat_sequence=merge_vars_with_tensor_arrays, |
| 2153 | expand_composites=True) |
| 2154 | c = ops.convert_to_tensor(pred(*packed_vars)) |
no test coverage detected