Repeatedly applies callable `fn` to a sequence of elements. Implemented by functional_ops.While, tpu friendly, no gradient. This is similar to functional_ops.scan but significantly faster on tpu/gpu for the forward backward use case. Examples: scan(lambda a, e: a + e, [1.0, 2.0, 3.0],
(fn, elems, initial, reverse=False, inclusive=False, final_only=False)
| 1112 | # functional_ops.scan, but could be replaced by that or something similar if |
| 1113 | # things change. |
| 1114 | def _scan(fn, elems, initial, reverse=False, inclusive=False, final_only=False): |
| 1115 | """Repeatedly applies callable `fn` to a sequence of elements. |
| 1116 | |
| 1117 | Implemented by functional_ops.While, tpu friendly, no gradient. |
| 1118 | |
| 1119 | This is similar to functional_ops.scan but significantly faster on tpu/gpu |
| 1120 | for the forward backward use case. |
| 1121 | |
| 1122 | Examples: |
| 1123 | scan(lambda a, e: a + e, [1.0, 2.0, 3.0], 1.0) => [2.0, 4.0, 7.0] |
| 1124 | |
| 1125 | Multiple accumulators: |
| 1126 | scan(lambda a, e: (a[0] + e, a[1] * e), [1.0, 2.0, 3.0], (0.0, 1.0)) |
| 1127 | |
| 1128 | Multiple inputs: |
| 1129 | scan(lambda a, e: a + (e[0] * e[1]), (elems1, elems2), 0.0) |
| 1130 | |
| 1131 | Args: |
| 1132 | fn: callable, fn(accumulators, element) return new accumulator values. The |
| 1133 | (possibly nested) sequence of accumulators is the same as `initial` and |
| 1134 | the return value must have the same structure. |
| 1135 | elems: A (possibly nested) tensor which will be unpacked along the first |
| 1136 | dimension. The resulting slices will be the second argument to fn. The |
| 1137 | first dimension of all nested input tensors must be the same. |
| 1138 | initial: A tensor or (possibly nested) sequence of tensors with initial |
| 1139 | values for the accumulators. |
| 1140 | reverse: (optional) True enables scan and output elems in reverse order. |
| 1141 | inclusive: (optional) True includes the initial accumulator values in the |
| 1142 | output. Length of output will be len(elem sequence) + 1. Not meaningful if |
| 1143 | final_only is True. |
| 1144 | final_only: (optional) When True, return only the final accumulated values, |
| 1145 | not the concatenation of accumulated values for each input. |
| 1146 | |
| 1147 | Returns: |
| 1148 | A (possibly nested) sequence of tensors with the results of applying fn |
| 1149 | to tensors unpacked from elems and previous accumulator values. |
| 1150 | """ |
| 1151 | |
| 1152 | flat_elems = [ops.convert_to_tensor(x) for x in nest.flatten(elems)] |
| 1153 | num_elems = array_ops.shape(flat_elems[0])[0] |
| 1154 | pack_elems = lambda x: nest.pack_sequence_as(structure=elems, flat_sequence=x) |
| 1155 | flat_initial = [ops.convert_to_tensor(x) for x in nest.flatten(initial)] |
| 1156 | pack = lambda x: nest.pack_sequence_as(structure=initial, flat_sequence=x) |
| 1157 | accum_dtypes = [x.dtype for x in flat_initial] |
| 1158 | num_accums = len(flat_initial) |
| 1159 | |
| 1160 | # Types for counter, [outputs], [accumulators] loop arguments. |
| 1161 | if final_only: |
| 1162 | loop_dtypes = [dtypes.int32, dtypes.int32] + accum_dtypes |
| 1163 | else: |
| 1164 | loop_dtypes = [dtypes.int32, dtypes.int32] + accum_dtypes + accum_dtypes |
| 1165 | |
| 1166 | # TODO(tombagby): Update to tfe.defun |
| 1167 | def cond(i, num_elems, *args): |
| 1168 | del args |
| 1169 | return i >= 0 if reverse else i < num_elems |
| 1170 | |
| 1171 | # The loop *args are [output tensors] + [accumulator tensors] which must |
no test coverage detected