Builds and returns the gradient FuncGraph of `func_graph` and its args. The returned grad_func_graph must be called with the returned args + grad_func_graph.captures. Args: ys: A `Tensor` or list of tensors to be differentiated. xs: A `Tensor` or list of tensors to be used for differ
(ys, xs, grads, cond_graph, body_graph, name, while_op,
maximum_iterations)
| 562 | |
| 563 | |
| 564 | def _create_grad_func(ys, xs, grads, cond_graph, body_graph, name, while_op, |
| 565 | maximum_iterations): |
| 566 | """Builds and returns the gradient FuncGraph of `func_graph` and its args. |
| 567 | |
| 568 | The returned grad_func_graph must be called with the returned |
| 569 | args + grad_func_graph.captures. |
| 570 | |
| 571 | Args: |
| 572 | ys: A `Tensor` or list of tensors to be differentiated. |
| 573 | xs: A `Tensor` or list of tensors to be used for differentiation. |
| 574 | grads: The incoming grads for `ys`. |
| 575 | cond_graph: FuncGraph for the forward cond function. |
| 576 | body_graph: FuncGraph for the forward body function. |
| 577 | name: Name of the returned gradient function. |
| 578 | while_op: The forward While op. |
| 579 | maximum_iterations: Tensor. The maximum number of iterations. |
| 580 | |
| 581 | Returns: |
| 582 | 2-tuple of (grad_func_graph, args). |
| 583 | """ |
| 584 | assert len(ys) == len(grads) |
| 585 | |
| 586 | total_iters = while_op.outputs[0] |
| 587 | counter = constant_op.constant( |
| 588 | 0, dtype=total_iters.dtype, name="grad_counter") |
| 589 | |
| 590 | # Build frozen sets so that we do not have linear time lookups in |
| 591 | # `_is_loop_invariant`. Note: `body_graph.inputs` and `body_graph.outputs` |
| 592 | # may get updated during gradient computation because we add accumulators to |
| 593 | # the forward op. However, those are not loop invariants so wouldn't affect |
| 594 | # the output of `_is_loop_invariant`. Also we would never attempt to capture |
| 595 | # those accumulators so `_is_loop_invariant` should never receive those new |
| 596 | # tensors as args. |
| 597 | body_graph_inputs = object_identity.ObjectIdentitySet(body_graph.inputs) |
| 598 | body_graph_outputs = object_identity.ObjectIdentitySet(body_graph.outputs) |
| 599 | |
| 600 | args = [counter, maximum_iterations, total_iters] + list(grads) |
| 601 | # Note: The returned function does not have `args` in the list of |
| 602 | # `external_captures`. |
| 603 | grad_func_graph = func_graph_module.func_graph_from_py_func( |
| 604 | name, |
| 605 | lambda *args: _grad_fn(ys, xs, args, body_graph), |
| 606 | args, {}, |
| 607 | func_graph=_WhileBodyGradFuncGraph(name, cond_graph, body_graph, |
| 608 | maximum_iterations, while_op, |
| 609 | body_graph_inputs, body_graph_outputs)) |
| 610 | |
| 611 | # Update the list of outputs with tensors corresponding to the captured |
| 612 | # tensors. We capture 3 types of tensors when building the grad fn: |
| 613 | # 1. Accumulators for forward graph intermediates which are not loop |
| 614 | # invariants. The outputs corresponding to these are populated in |
| 615 | # `popped_tensor_lists` by `_WhileBodyGradFuncGraph`. |
| 616 | # 2. Resources, which are output as is. |
| 617 | # 3. Forward graph loop invariants, which are output as is. |
| 618 | for external_capture, internal_capture in grad_func_graph.captures: |
| 619 | if ops.tensor_id(internal_capture) in grad_func_graph.popped_tensor_lists: |
| 620 | new_output = grad_func_graph.popped_tensor_lists[ops.tensor_id( |
| 621 | internal_capture)] |
no test coverage detected