Add the loop termination condition and body to the graph.
(self, pred, body, loop_vars, shape_invariants,
return_same_structure)
| 2220 | return original_body_result, exit_vars |
| 2221 | |
| 2222 | def BuildLoop(self, pred, body, loop_vars, shape_invariants, |
| 2223 | return_same_structure): |
| 2224 | """Add the loop termination condition and body to the graph.""" |
| 2225 | |
| 2226 | # Keep original_loop_vars to identify which are TensorArrays |
| 2227 | original_loop_vars = loop_vars |
| 2228 | # Convert TensorArrays to their flow variables |
| 2229 | loop_vars = nest.map_structure( |
| 2230 | _convert_tensorarray_to_flow, |
| 2231 | nest.flatten(loop_vars, expand_composites=False), |
| 2232 | expand_composites=True) |
| 2233 | loop_vars = ops.convert_n_to_tensor_or_composite(loop_vars) |
| 2234 | if shape_invariants is None: |
| 2235 | shape_invariants = nest.map_structure( |
| 2236 | _get_shape_invariant, loop_vars, expand_composites=False) |
| 2237 | loop_vars = nest.flatten(loop_vars, expand_composites=True) |
| 2238 | try: |
| 2239 | self.Enter() |
| 2240 | # _BuildLoop calls _update_input in several places. _mutation_lock() |
| 2241 | # ensures a Session.run call cannot occur between creating and mutating |
| 2242 | # new ops. |
| 2243 | with ops.get_default_graph()._mutation_lock(): # pylint: disable=protected-access |
| 2244 | original_body_result, exit_vars = self._BuildLoop( |
| 2245 | pred, body, original_loop_vars, loop_vars, shape_invariants) |
| 2246 | finally: |
| 2247 | self.Exit() |
| 2248 | |
| 2249 | flat_result = nest.flatten(original_body_result, expand_composites=True) |
| 2250 | # Convert TensorArray flow variables outside the context back into |
| 2251 | # their associated TensorArrays for returning to caller. |
| 2252 | exit_vars_with_tensor_arrays = ( |
| 2253 | _convert_flows_to_tensorarrays(flat_result, exit_vars)) |
| 2254 | packed_exit_vars = nest.pack_sequence_as( |
| 2255 | structure=original_body_result, |
| 2256 | flat_sequence=exit_vars_with_tensor_arrays, |
| 2257 | expand_composites=True) |
| 2258 | |
| 2259 | if return_same_structure: |
| 2260 | return packed_exit_vars |
| 2261 | else: |
| 2262 | return packed_exit_vars[0] if len(exit_vars) == 1 else packed_exit_vars |
| 2263 | |
| 2264 | def _FixControlInputsAndContext(self, enters): |
| 2265 | graph = ops.get_default_graph() |
no test coverage detected