Returns `FuncGraph`s for the input op branches. Args: op: The If or Case Operation. Returns: A tuple of the `FuncGraph`s of the then_branch and else_branch (all branches for Case).
(op)
| 296 | |
| 297 | |
| 298 | def get_func_graphs(op): |
| 299 | """Returns `FuncGraph`s for the input op branches. |
| 300 | |
| 301 | Args: |
| 302 | op: The If or Case Operation. |
| 303 | |
| 304 | Returns: |
| 305 | A tuple of the `FuncGraph`s of the then_branch and else_branch (all branches |
| 306 | for Case). |
| 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")), |
| 332 | _get_func_graph_for_branch(op.get_attr("else_branch"))) |
| 333 | elif op.type == "Case": |
| 334 | return [_get_func_graph_for_branch(branch_fn) |
| 335 | for branch_fn in op.get_attr("branches")] |
| 336 | else: |
| 337 | raise ValueError("Unsupported op type: {}".format(op.type)) |
| 338 | |
| 339 | |
| 340 | def _grad_fn(func_graph, grads): |
no test coverage detected