Convert the body function.
(self, pfor_input, inputs_stacked,
new_indices, cond_stacked, new_inputs,
not_all_done)
| 500 | return not_all_done, new_indices, new_inputs, new_output_tas |
| 501 | |
| 502 | def _process_body(self, pfor_input, inputs_stacked, |
| 503 | new_indices, cond_stacked, new_inputs, |
| 504 | not_all_done): |
| 505 | """Convert the body function.""" |
| 506 | |
| 507 | def true_fn(control_inputs, body_pfor, body_output, stacked): |
| 508 | """Converts the body function for all but last iteration. |
| 509 | |
| 510 | This essentially converts body_output. Additionally, it needs to handle |
| 511 | any control dependencies on the NextIteration node. So it creates another |
| 512 | Identity node with the converted dependencies. |
| 513 | """ |
| 514 | converted_control_inp = [] |
| 515 | for x in control_inputs: |
| 516 | for t in x.outputs: |
| 517 | converted_control_inp.append(body_pfor._convert_helper(t).t) |
| 518 | if stacked: |
| 519 | # Note convert always does the stacking. |
| 520 | output = body_pfor.convert(body_output) |
| 521 | else: |
| 522 | output, convert_stacked, _ = body_pfor._convert_helper(body_output) |
| 523 | assert convert_stacked == stacked, body_output |
| 524 | with ops.control_dependencies(converted_control_inp): |
| 525 | return array_ops.identity(output) |
| 526 | |
| 527 | body_pfor = self._init_pfor(pfor_input.pfor, new_indices, |
| 528 | cond_stacked, new_inputs, |
| 529 | inputs_stacked) |
| 530 | new_outputs = [] |
| 531 | |
| 532 | for i, (body_output, stacked) in enumerate( |
| 533 | zip(self._body_outputs, inputs_stacked)): |
| 534 | control_inp = self._next_iter_control_inputs[i] |
| 535 | out_dtype = body_output.dtype |
| 536 | # Note that we want to run the body only if not all pfor iterations are |
| 537 | # done. If all are done, we return empty tensors since these values will |
| 538 | # not be used. Notice that the value returned by the loop is based on |
| 539 | # TensorArrays and not directly on these returned values. |
| 540 | # pylint: disable=cell-var-from-loop |
| 541 | new_output = control_flow_ops.cond( |
| 542 | not_all_done, |
| 543 | lambda: true_fn(control_inp, body_pfor, body_output, stacked), |
| 544 | lambda: constant_op.constant([], dtype=out_dtype)) |
| 545 | # pylint: enable=cell-var-from-loop |
| 546 | new_outputs.append(new_output) |
| 547 | return new_outputs |
| 548 | |
| 549 | def __call__(self, pfor_input): |
| 550 | """Converter for the while_loop. |