Returns the values that should be returned from the while grad function. Args: outputs: the raw Tensor outputs of the grad While op. grads: the input gradients to the gradient function. body_grad_graph: _WhileBodyGradFuncGraph. Returns: A list of gradient values. May include No
(outputs, grads, body_grad_graph)
| 719 | |
| 720 | |
| 721 | def _get_structured_grad_output(outputs, grads, body_grad_graph): |
| 722 | """Returns the values that should be returned from the while grad function. |
| 723 | |
| 724 | Args: |
| 725 | outputs: the raw Tensor outputs of the grad While op. |
| 726 | grads: the input gradients to the gradient function. |
| 727 | body_grad_graph: _WhileBodyGradFuncGraph. |
| 728 | |
| 729 | Returns: |
| 730 | A list of gradient values. May include Nones. |
| 731 | """ |
| 732 | result = [] |
| 733 | # outputs[0] is the loop counter. |
| 734 | # outputs[1] is maximum_iterations. |
| 735 | # outputs[2] is the total number of loop iterations. |
| 736 | outputs_idx = 3 |
| 737 | structured_outputs_idx = 3 |
| 738 | for g in grads: |
| 739 | # Set None as the output gradient for tensors with None input gradient. |
| 740 | if g is None: |
| 741 | result.append(None) |
| 742 | continue |
| 743 | output = body_grad_graph.structured_outputs[structured_outputs_idx] |
| 744 | structured_outputs_idx += 1 |
| 745 | if isinstance(output, ops.IndexedSlices): |
| 746 | # TODO(skyewm): is there a more robust way to determine the order of |
| 747 | # flattened IndexedSlices components? |
| 748 | result.append(ops.IndexedSlices( |
| 749 | values=outputs[outputs_idx], |
| 750 | indices=outputs[outputs_idx + 1], |
| 751 | dense_shape=outputs[outputs_idx + 2])) |
| 752 | outputs_idx += 3 |
| 753 | else: |
| 754 | assert isinstance(output, ops.Tensor) |
| 755 | result.append(outputs[outputs_idx]) |
| 756 | outputs_idx += 1 |
| 757 | |
| 758 | return result |
| 759 | |
| 760 | |
| 761 | def _get_accumulator(tensor): |