FuncGraph for the gradient function of the body of a While op. Contains the logic for capturing the tensors from the body of the forward While op which is as follows: 1. If the tensor is of resource type (these are not accumulated): a. Ensure that the tensor is a loop invariant, i.e., it
| 828 | |
| 829 | |
| 830 | class _WhileBodyGradFuncGraph(util.WhileBodyFuncGraph): |
| 831 | """FuncGraph for the gradient function of the body of a While op. |
| 832 | |
| 833 | Contains the logic for capturing the tensors from the body of the forward |
| 834 | While op which is as follows: |
| 835 | 1. If the tensor is of resource type (these are not accumulated): |
| 836 | a. Ensure that the tensor is a loop invariant, i.e., it exists in both loop |
| 837 | inputs and outputs at the same index. |
| 838 | b. Lookup the corresponding resource tensor in the forward outer graph and |
| 839 | try to capture that. |
| 840 | 2. If the tensor is not of resource type: |
| 841 | a. Create an accumulator for that tensor and output it from the forward |
| 842 | pass. Note this also requires adding it as an input to the forward pass. |
| 843 | b. Capture the accumulator from the forward pass in this FuncGraph. This |
| 844 | will later be resolved to the correct output of the forward While op. |
| 845 | c. Pop a value from the captured placeholder and use it as the captured |
| 846 | value for the forward pass tensor. |
| 847 | |
| 848 | This only allows capturing tensors in the forward graph. A ValueError is |
| 849 | raised if an attempt is made to capture a tensor not in the forward graph. |
| 850 | To manually capture capture a tensor that is not in the forward graph, call |
| 851 | `capture` with `whitelisted=True`. |
| 852 | |
| 853 | Note: The `captures` dict does not contain the forward tensor since it is not |
| 854 | directly captured. It contains the accumulator corresponding to this forward |
| 855 | tensor. |
| 856 | |
| 857 | Attributes: |
| 858 | while_op_needs_rewrite: True if any non-resource intermediates were |
| 859 | captured, meaning the forward While op needs to be rewritten to output the |
| 860 | corresponding accumulators. |
| 861 | empty_tensor_lists: list of EmptyTensorList tensors to be used as initial |
| 862 | input to the new accumulators in the forward graph. |
| 863 | popped_tensor_lists: dict from the captured accumulator placeholder to the |
| 864 | TensorList obtained after popping the intermediate tensor from it. The |
| 865 | values of this dict need to be added to the list of outputs. |
| 866 | """ |
| 867 | |
| 868 | def __init__(self, name, forward_cond_graph, forward_body_graph, |
| 869 | maximum_iterations, forward_while_op, body_graph_inputs, |
| 870 | body_graph_outputs): |
| 871 | super(_WhileBodyGradFuncGraph, self).__init__(name) |
| 872 | self.empty_tensor_lists = [] |
| 873 | self.popped_tensor_lists = {} |
| 874 | # FuncGraph for the body of the forward While op. |
| 875 | self._forward_graph = forward_body_graph |
| 876 | # FuncGraph for the cond of the forward While op. |
| 877 | self._forward_cond_graph = forward_cond_graph |
| 878 | self._maximum_iterations = maximum_iterations |
| 879 | self._forward_while_op = forward_while_op |
| 880 | # Only for use in `_is_loop_invariant`. These are not updated when |
| 881 | # additional tensors are added to `forward_body_graph.inputs` and |
| 882 | # `forward_body_graph.outputs` in `_capture_helper`. |
| 883 | self._forward_graph_inputs = body_graph_inputs |
| 884 | self._forward_graph_outputs = body_graph_outputs |
| 885 | # Dict from forward intermediate tensor to its indirectly captured tensor |
| 886 | # in this graph. Indirect capturing happens in two ways: |
| 887 | # 1. For non-resource tensors we capture their accumulators from the forward |