Returns the tensors to pass as captured inputs to `body_grad_graph`. `body_grad_graph` may have external references to: 1. Its outer graph containing the input gradients. These are left as-is. 2. Accumulators captured from the forward-pass graph. These should have been added as `while_op
(body_graph, body_grad_graph, while_op)
| 674 | |
| 675 | |
| 676 | def _resolve_grad_captures(body_graph, body_grad_graph, while_op): |
| 677 | """Returns the tensors to pass as captured inputs to `body_grad_graph`. |
| 678 | |
| 679 | `body_grad_graph` may have external references to: |
| 680 | 1. Its outer graph containing the input gradients. These are left as-is. |
| 681 | 2. Accumulators captured from the forward-pass graph. These should have been |
| 682 | added as `while_op` outputs after the gradient graph was built. We replace |
| 683 | these with the corresponding output of `while_op`, i.e. a tensor in |
| 684 | `body_graph.outer_graph`. In the case of nested control flow or functions, |
| 685 | the gradient logic handling `body_grad_graph.outer_graph` will make sure |
| 686 | the tensor from `body_graph.outer_graph` is also correctly captured. |
| 687 | |
| 688 | Args: |
| 689 | body_graph: FuncGraph. The forward-pass body function. |
| 690 | body_grad_graph: FuncGraph. The body gradients function. |
| 691 | while_op: The forward-pass While Operation calling `body_graph`. |
| 692 | |
| 693 | Returns: |
| 694 | A list of input tensors to be passed as the captured inputs to |
| 695 | `body_grad_graph`. |
| 696 | """ |
| 697 | new_capture_inputs = [] |
| 698 | for t in body_grad_graph.external_captures: |
| 699 | # All values captured by gradient computation should be from the forward |
| 700 | # graph or a captured resource variable (note that input gradients are |
| 701 | # regular non-captured inputs). |
| 702 | if t.graph == body_graph: |
| 703 | # Captured accumulator or loop invariant. |
| 704 | for i, output in enumerate(t.graph.outputs): |
| 705 | if output is t: |
| 706 | t = while_op.outputs[i] |
| 707 | break |
| 708 | |
| 709 | # Note: We rely on the capturing logic of the gradient While op graph to |
| 710 | # correctly capture the tensors in `body_graph.outer_graph`. Both cond_v2 |
| 711 | # and while_v2 handle this while building their gradient functions. |
| 712 | assert t.graph == body_graph.outer_graph |
| 713 | else: |
| 714 | # Captured resource variable |
| 715 | assert t.dtype == dtypes.resource |
| 716 | |
| 717 | new_capture_inputs.append(t) |
| 718 | return new_capture_inputs |
| 719 | |
| 720 | |
| 721 | def _get_structured_grad_output(outputs, grads, body_grad_graph): |