Runs `loop_fn` `iters` times and stacks the outputs. Runs `loop_fn` `iters` times, with input values from 0 to `iters - 1`, and stacks corresponding outputs of the different runs. Args: loop_fn: A function that takes an int32 scalar tf.Tensor object representing the iteration numb
(loop_fn, loop_fn_dtypes, iters, parallel_iterations=None)
| 42 | |
| 43 | |
| 44 | def for_loop(loop_fn, loop_fn_dtypes, iters, parallel_iterations=None): |
| 45 | """Runs `loop_fn` `iters` times and stacks the outputs. |
| 46 | |
| 47 | |
| 48 | Runs `loop_fn` `iters` times, with input values from 0 to `iters - 1`, and |
| 49 | stacks corresponding outputs of the different runs. |
| 50 | |
| 51 | Args: |
| 52 | loop_fn: A function that takes an int32 scalar tf.Tensor object representing |
| 53 | the iteration number, and returns a possibly nested structure of tensor |
| 54 | objects. The shape of these outputs should not depend on the input. |
| 55 | loop_fn_dtypes: dtypes for the outputs of loop_fn. |
| 56 | iters: Number of iterations for which to run loop_fn. |
| 57 | parallel_iterations: The number of iterations that can be dispatched in |
| 58 | parallel. This knob can be used to control the total memory usage. |
| 59 | |
| 60 | Returns: |
| 61 | Returns a nested structure of stacked output tensor objects with the same |
| 62 | nested structure as the output of `loop_fn`. |
| 63 | """ |
| 64 | |
| 65 | flat_loop_fn_dtypes = nest.flatten(loop_fn_dtypes) |
| 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} |
| 89 | else: |
| 90 | extra_args = {} |
| 91 | ta_list = control_flow_ops.while_loop( |
| 92 | lambda i, *ta: i < iters, |
| 93 | while_body, |
| 94 | [0] + [tensor_array_ops.TensorArray(dtype.base_dtype, iters) |
| 95 | for dtype in flat_loop_fn_dtypes], |
| 96 | **extra_args)[1:] |
| 97 | |
| 98 | # TODO(rachelim): enable this for sparse tensors |
| 99 | |
| 100 | output = [None if is_none else ta.concat() |
| 101 | for ta, is_none in zip(ta_list, is_none_list)] |
no test coverage detected