Like tf.while_loop, except emits a single While op.
(cond,
body,
loop_vars,
shape_invariants=None,
parallel_iterations=10,
maximum_iterations=None,
name=None,
return_same_structure=True,
back_prop=True)
| 58 | |
| 59 | |
| 60 | def while_loop(cond, |
| 61 | body, |
| 62 | loop_vars, |
| 63 | shape_invariants=None, |
| 64 | parallel_iterations=10, |
| 65 | maximum_iterations=None, |
| 66 | name=None, |
| 67 | return_same_structure=True, |
| 68 | back_prop=True): |
| 69 | """Like tf.while_loop, except emits a single While op.""" |
| 70 | # Keep the original loop_vars around to know which args were TensorArrays. |
| 71 | orig_loop_vars = loop_vars |
| 72 | # Cache its length since we use it at multiple places below. |
| 73 | len_orig_loop_vars = len(orig_loop_vars) |
| 74 | |
| 75 | # Convert TensorArrays to their flow variables. These get converted back to |
| 76 | # TensorArrays before calling `cond` and `body`. See `wrapped_cond` and |
| 77 | # `wrapped_body` below. |
| 78 | loop_vars = list(_tensor_array_to_flow(orig_loop_vars)) |
| 79 | loop_vars = nest.map_structure( |
| 80 | ops.internal_convert_to_tensor_or_indexed_slices, loop_vars, |
| 81 | expand_composites=True) |
| 82 | if shape_invariants is not None: |
| 83 | nest.assert_same_structure(orig_loop_vars, shape_invariants, |
| 84 | expand_composites=False) |
| 85 | signature = nest.map_structure( |
| 86 | control_flow_ops._shape_invariant_to_type_spec, loop_vars, |
| 87 | list(shape_invariants), expand_composites=False) |
| 88 | shape_invariants = nest.map_structure( |
| 89 | control_flow_ops._get_shape_invariant, loop_vars, |
| 90 | list(shape_invariants), expand_composites=False) |
| 91 | |
| 92 | else: |
| 93 | signature = nest.map_structure( |
| 94 | type_spec.type_spec_from_value, loop_vars, expand_composites=False) |
| 95 | shape_invariants = nest.map_structure( |
| 96 | control_flow_ops._get_shape_invariant, loop_vars, |
| 97 | expand_composites=False) |
| 98 | if not name: |
| 99 | name = "while" |
| 100 | |
| 101 | with ops.name_scope(name) as scope: |
| 102 | with ops.name_scope(None): |
| 103 | cond_name = util.unique_fn_name(scope, "cond") |
| 104 | body_name = util.unique_fn_name(scope, "body") |
| 105 | maximum_iterations_loop_var = _build_maximum_iterations_loop_var( |
| 106 | maximum_iterations) |
| 107 | loop_counter = constant_op.constant( |
| 108 | 0, |
| 109 | dtype=maximum_iterations_loop_var.dtype |
| 110 | if maximum_iterations is not None else None, |
| 111 | name="loop_counter") |
| 112 | # Add loop counter needed for computing gradients. |
| 113 | loop_vars = [loop_counter, maximum_iterations_loop_var] + loop_vars |
| 114 | |
| 115 | shape_invariants = [tensor_shape.TensorShape([])] * 2 + shape_invariants |
| 116 | signature = ( |
| 117 | [tensor_spec.TensorSpec.from_tensor(loop_counter), |