Body of while loop.
(i, *ta_list)
| 66 | is_none_list = [] |
| 67 | |
| 68 | def while_body(i, *ta_list): |
| 69 | """Body of while loop.""" |
| 70 | fn_output = nest.flatten(loop_fn(i)) |
| 71 | if len(fn_output) != len(flat_loop_fn_dtypes): |
| 72 | raise ValueError( |
| 73 | "Number of expected outputs, %d, does not match the number of " |
| 74 | "actual outputs, %d, from loop_fn" % (len(flat_loop_fn_dtypes), |
| 75 | len(fn_output))) |
| 76 | outputs = [] |
| 77 | del is_none_list[:] |
| 78 | is_none_list.extend([x is None for x in fn_output]) |
| 79 | for out, ta in zip(fn_output, ta_list): |
| 80 | # TODO(agarwal): support returning Operation objects from loop_fn. |
| 81 | if out is not None: |
| 82 | # out may be a ref tensor, wrap it in identity to get a non-ref tensor. |
| 83 | ta = ta.write(i, array_ops.expand_dims(out, 0)) |
| 84 | outputs.append(ta) |
| 85 | return tuple([i + 1] + outputs) |
| 86 | |
| 87 | if parallel_iterations is not None: |
| 88 | extra_args = {"parallel_iterations": parallel_iterations} |