Returns `FuncGraph` for the given function attribute. Args: while_op: The While Operation. func_attr_name: string Returns: `FuncGraph`
(while_op, func_attr_name)
| 535 | |
| 536 | # TODO(srbs): Pull this into common utils for cond_v2 and while_v2. |
| 537 | def _get_graph(while_op, func_attr_name): |
| 538 | """Returns `FuncGraph` for the given function attribute. |
| 539 | |
| 540 | Args: |
| 541 | while_op: The While Operation. |
| 542 | func_attr_name: string |
| 543 | |
| 544 | Returns: |
| 545 | `FuncGraph` |
| 546 | """ |
| 547 | # TODO(srbs): Handle TensorShapeProto in function_def_to_graph.input_shapes. |
| 548 | input_shapes = [ |
| 549 | tensor_shape.TensorShape(s) for s in while_op.get_attr("output_shapes") |
| 550 | ] |
| 551 | func_name = while_op.get_attr(func_attr_name).name |
| 552 | fdef = while_op.graph._get_function(func_name).definition |
| 553 | # `while_op.graph` may not be the same as `ops.get_default_graph()` e.g. |
| 554 | # if the `while_op` is in the body of another if/while/defun. We build the |
| 555 | # `func_graph` with `while_op.graph` as its `outer_graph`. This resembles how |
| 556 | # the `FuncGraph` was built in the forward pass. We need this so that we can |
| 557 | # appropriately capture references to outer tensors in the nested grad graphs. |
| 558 | with while_op.graph.as_default(): |
| 559 | func_graph = function_def_to_graph.function_def_to_graph(fdef, input_shapes) |
| 560 | func_graph._while = while_op |
| 561 | return func_graph |
| 562 | |
| 563 | |
| 564 | def _create_grad_func(ys, xs, grads, cond_graph, body_graph, name, while_op, |
no test coverage detected