Generates and returns a FuncGraph for the given branch.
(name_attr_list)
| 307 | """ |
| 308 | |
| 309 | def _get_func_graph_for_branch(name_attr_list): |
| 310 | """Generates and returns a FuncGraph for the given branch.""" |
| 311 | inputs = op.inputs[1:] # First input is pred. |
| 312 | input_shapes = [t.shape for t in inputs] |
| 313 | fdef = op.graph._get_function(name_attr_list.name).definition |
| 314 | # `op.graph` may not be the same as `ops.get_default_graph()` e.g. |
| 315 | # in the case of nested if ops or when the gradient is being computed |
| 316 | # from inside a Defun. We build the `func_graph` with `op.graph` as its |
| 317 | # `outer_graph`. This resembles how the `FuncGraph` was built in the |
| 318 | # forward pass. We need this so that we can resolve references to tensors |
| 319 | # in `func_graph` from its gradient graph in `_resolve_grad_inputs`. |
| 320 | with op.graph.as_default(): |
| 321 | func_graph = function_def_to_graph.function_def_to_graph( |
| 322 | fdef, input_shapes) |
| 323 | for external_t, internal_t in zip(inputs, func_graph.inputs): |
| 324 | custom_gradient.copy_handle_data(external_t, internal_t) |
| 325 | func_graph.reset_captures(zip(inputs, func_graph.inputs)) |
| 326 | # Link the op so that the gradient code can use it. |
| 327 | func_graph._forward_cond = op |
| 328 | return func_graph |
| 329 | |
| 330 | if op.type in ["If", "StatelessIf"]: |
| 331 | return (_get_func_graph_for_branch(op.get_attr("then_branch")), |
no test coverage detected