The gradient of an If op produced by cond_v2.
(op, *grads)
| 104 | @ops.RegisterGradient("StatelessIf") |
| 105 | @ops.RegisterGradient("If") |
| 106 | def _IfGrad(op, *grads): # pylint: disable=invalid-name |
| 107 | """The gradient of an If op produced by cond_v2.""" |
| 108 | # Get the if operator (this logic handles the case where op is a MockOp) |
| 109 | if_op = op.outputs[0].op |
| 110 | true_graph, false_graph = get_func_graphs(if_op) |
| 111 | # Note: op.graph != ops.get_default_graph() when we are computing the gradient |
| 112 | # of a nested cond. |
| 113 | assert true_graph.outer_graph == if_op.graph |
| 114 | assert false_graph.outer_graph == if_op.graph |
| 115 | |
| 116 | # Create grad functions that compute the gradient of the true/false forward |
| 117 | # graphs. These functions will capture tensors from the forward pass |
| 118 | # functions. |
| 119 | true_grad_graph = _create_grad_func( |
| 120 | true_graph, grads, util.unique_grad_fn_name(true_graph.name)) |
| 121 | false_grad_graph = _create_grad_func( |
| 122 | false_graph, grads, util.unique_grad_fn_name(false_graph.name)) |
| 123 | |
| 124 | if (true_grad_graph.op_needs_rewrite or false_grad_graph.op_needs_rewrite): |
| 125 | # Modify 'op' to output the intermediates needed by the grad functions. Note |
| 126 | # that all needed intermediates are wrapped in optionals. Each optional |
| 127 | # intermediate output will have a value iff its corresponding branch is |
| 128 | # taken. |
| 129 | # NOTE(skyewm): if there are any active sessions, this modification to `op` |
| 130 | # may make them unrunnable! |
| 131 | |
| 132 | if control_flow_util.GraphOrParentsInXlaContext(ops.get_default_graph()): |
| 133 | # XLA does not yet support optionals, so output intermediates directly and |
| 134 | # make them match via FakeParams, which can be converted to zeros in XLA. |
| 135 | # TODO(skyewm,jpienaar): can XLA support optionals? |
| 136 | true_intermediates = true_grad_graph.xla_intermediates |
| 137 | false_intermediates = false_grad_graph.xla_intermediates |
| 138 | extra_true_outputs, extra_false_outputs = _make_intermediates_match_xla( |
| 139 | [true_graph, false_graph], [true_intermediates, false_intermediates]) |
| 140 | else: |
| 141 | true_intermediates = true_grad_graph.wrapped_intermediates |
| 142 | false_intermediates = false_grad_graph.wrapped_intermediates |
| 143 | # Make outputs match by adding none optionals. |
| 144 | extra_true_outputs, extra_false_outputs = _make_intermediates_match( |
| 145 | [true_graph, false_graph], [true_intermediates, false_intermediates]) |
| 146 | |
| 147 | true_graph.outputs.extend(extra_true_outputs) |
| 148 | false_graph.outputs.extend(extra_false_outputs) |
| 149 | # TODO(skyewm): indicate it's an internal bug if this fails. |
| 150 | _check_same_outputs(_COND, [true_graph, false_graph]) |
| 151 | |
| 152 | true_graph.name += "_rewritten" |
| 153 | false_graph.name += "_rewritten" |
| 154 | |
| 155 | if_op._set_func_attr("then_branch", util.create_new_tf_function(true_graph)) |
| 156 | if_op._set_func_attr("else_branch", |
| 157 | util.create_new_tf_function(false_graph)) |
| 158 | if_op._set_type_list_attr("Tout", true_graph.output_types) |
| 159 | if_op._set_shape_list_attr("output_shapes", true_graph.output_shapes) |
| 160 | if_op._add_outputs( |
| 161 | [t.dtype for t in extra_true_outputs], |
| 162 | [t.shape for t in extra_true_outputs]) |
| 163 |
nothing calls this directly
no test coverage detected