Loop body.
(i, num_elems, *args)
| 1171 | # The loop *args are [output tensors] + [accumulator tensors] which must |
| 1172 | # be paired. Each output corresponds to one accumulator. |
| 1173 | def body(i, num_elems, *args): |
| 1174 | """Loop body.""" |
| 1175 | i.set_shape([]) |
| 1176 | if final_only: |
| 1177 | accum = args |
| 1178 | else: |
| 1179 | out, accum = args[:num_accums], args[num_accums:] |
| 1180 | slices = [array_ops.gather(e, i) for e in flat_elems] |
| 1181 | accum = fn(pack(accum), pack_elems(slices)) |
| 1182 | flat_accum = nest.flatten(accum) |
| 1183 | if final_only: |
| 1184 | new_out = [] |
| 1185 | else: |
| 1186 | update_i = i + 1 if inclusive and not reverse else i |
| 1187 | new_out = [ |
| 1188 | inplace_ops.alias_inplace_update(x, update_i, y) |
| 1189 | for x, y in zip(out, flat_accum) |
| 1190 | ] |
| 1191 | i = i - 1 if reverse else i + 1 |
| 1192 | return [i, num_elems] + new_out + flat_accum |
| 1193 | |
| 1194 | init_i = ( |
| 1195 | array_ops.shape(flat_elems[0])[0] - |