Extra `cond` wrapper that can handle the extra counter loop_var.
(loop_counter, maximum_iterations_arg, *args)
| 123 | add_control_dependencies = ops.get_default_graph()._add_control_dependencies |
| 124 | |
| 125 | def wrapped_cond(loop_counter, maximum_iterations_arg, *args): |
| 126 | """Extra `cond` wrapper that can handle the extra counter loop_var.""" |
| 127 | # Convert the flow variables in `args` to TensorArrays. `args` should |
| 128 | # already have the same structure as `orig_loop_vars` but currently there |
| 129 | # is no nest.zip so we call `_pack_sequence_as` which flattens both |
| 130 | # `orig_loop_vars` and `args`, converts flows in `args` to TensorArrays |
| 131 | # and packs it into the structure of `orig_loop_vars`. |
| 132 | pred = cond(*_pack_sequence_as(orig_loop_vars, args)) |
| 133 | if (tensor_util.is_tensor(pred) and |
| 134 | (pred.shape.dims is None or pred.shape.dims)): |
| 135 | pred = array_ops.squeeze_v2(pred) |
| 136 | |
| 137 | if maximum_iterations is None: |
| 138 | return pred |
| 139 | else: |
| 140 | return math_ops.logical_and( |
| 141 | loop_counter < maximum_iterations_arg, pred) |
| 142 | |
| 143 | # NOTE(skyewm): we set collections to the outer graph's collections for |
| 144 | # compatibility with TPUEstimator. |
nothing calls this directly
no test coverage detected