The loop body of map_fn. Args: i: the loop counter tas: the flat TensorArray accumulator list Returns: (i + 1, tas): the updated counter + updated TensorArrays Raises: TypeError: if dtype and packed_fn_values structure do not match ValueTy
(i, tas)
| 240 | for dt in dtype_flat] |
| 241 | |
| 242 | def compute(i, tas): |
| 243 | """The loop body of map_fn. |
| 244 | |
| 245 | Args: |
| 246 | i: the loop counter |
| 247 | tas: the flat TensorArray accumulator list |
| 248 | |
| 249 | Returns: |
| 250 | (i + 1, tas): the updated counter + updated TensorArrays |
| 251 | |
| 252 | Raises: |
| 253 | TypeError: if dtype and packed_fn_values structure do not match |
| 254 | ValueType: if dtype and packed_fn_values lengths do not match |
| 255 | """ |
| 256 | packed_values = input_pack([elem_ta.read(i) for elem_ta in elems_ta]) |
| 257 | packed_fn_values = fn(packed_values) |
| 258 | nest.assert_same_structure(dtype or elems, packed_fn_values) |
| 259 | flat_fn_values = output_flatten(packed_fn_values) |
| 260 | tas = [ta.write(i, value) for (ta, value) in zip(tas, flat_fn_values)] |
| 261 | return (i + 1, tas) |
| 262 | |
| 263 | _, r_a = control_flow_ops.while_loop( |
| 264 | lambda i, _: i < n, compute, (i, accs_ta), |