Returns the tensors to pass as inputs to `grad_graph`. The `grad_graph` may have external references to 1. Its outer graph containing the input gradients. These references are kept as is. 2. Tensors in the forward pass graph. These tensors may not be "live" when the gradient is bein
(cond_graph, grad_graph)
| 396 | |
| 397 | |
| 398 | def _resolve_grad_inputs(cond_graph, grad_graph): |
| 399 | """Returns the tensors to pass as inputs to `grad_graph`. |
| 400 | |
| 401 | The `grad_graph` may have external references to |
| 402 | 1. Its outer graph containing the input gradients. These references are kept |
| 403 | as is. |
| 404 | 2. Tensors in the forward pass graph. These tensors may not be "live" |
| 405 | when the gradient is being computed. We replace such references by their |
| 406 | corresponding tensor in `cond_graph.outer_graph`. In the case of nested |
| 407 | control flow or functions, the gradient logic handling |
| 408 | `grad_graph.outer_graph` will make sure the tensor from |
| 409 | `cond_graph.outer_graph` is also correctly captured. |
| 410 | |
| 411 | Args: |
| 412 | cond_graph: FuncGraph. The forward-pass function. |
| 413 | grad_graph: FuncGraph. The gradients function. |
| 414 | |
| 415 | Returns: |
| 416 | A list of inputs tensors to be passed to grad_graph. |
| 417 | """ |
| 418 | new_inputs = [] |
| 419 | |
| 420 | for t in grad_graph.external_captures: |
| 421 | # `t` must either be in `grad_graph.outer_graph` or in the forward |
| 422 | # `cond_graph`. |
| 423 | if t.graph != grad_graph.outer_graph: |
| 424 | assert t.graph == cond_graph |
| 425 | # `internal_captures` are not treated as intermediates and hence not added |
| 426 | # to If op outputs. So we get the outer tensor corresponding to those |
| 427 | # from the list of `external_captures`. |
| 428 | for i, output in enumerate(t.graph.outputs): |
| 429 | if output is t: |
| 430 | t = t.graph._forward_cond.outputs[i] |
| 431 | break |
| 432 | else: |
| 433 | for i, output in enumerate(t.graph.internal_captures): |
| 434 | if output is t: |
| 435 | t = t.graph.external_captures[i] |
| 436 | break |
| 437 | else: |
| 438 | raise ValueError("Could not find external tensor capture {tensor} in " |
| 439 | "captures or outputs".format(tensor=t)) |
| 440 | |
| 441 | # Note: We rely on the capturing logic of the gradient If op graph to |
| 442 | # correctly capture the tensors in `cond_graph.outer_graph`. Both cond_v2 |
| 443 | # and while_v2 handle this while building their gradient functions. |
| 444 | assert t.graph == cond_graph.outer_graph |
| 445 | new_inputs.append(t) |
| 446 | |
| 447 | return new_inputs |
| 448 | |
| 449 | |
| 450 | def _get_intermediates(func_graph): |