Gradients for a Merge op are calculated using a Switch op.
(op, grad, _)
| 94 | |
| 95 | @ops.RegisterGradient("Merge") |
| 96 | def _MergeGrad(op, grad, _): |
| 97 | """Gradients for a Merge op are calculated using a Switch op.""" |
| 98 | input_op = op.inputs[0].op |
| 99 | graph = ops.get_default_graph() |
| 100 | # pylint: disable=protected-access |
| 101 | op_ctxt = control_flow_util.GetOutputContext(input_op) |
| 102 | grad_ctxt = graph._get_control_flow_context() |
| 103 | # pylint: enable=protected-access |
| 104 | if isinstance(op_ctxt, WhileContext): |
| 105 | # pylint: disable=protected-access |
| 106 | return control_flow_ops._SwitchRefOrTensor(grad, grad_ctxt.pivot) |
| 107 | # pylint: enable=protected-access |
| 108 | elif isinstance(op_ctxt, CondContext): |
| 109 | pred = op_ctxt.pred |
| 110 | if grad_ctxt and grad_ctxt.grad_state: |
| 111 | # This Merge node is part of a cond within a loop. |
| 112 | # The backprop needs to have the value of this predicate for every |
| 113 | # iteration. So we must have its values accumulated in the forward, and |
| 114 | # use the accumulated values as the predicate for this backprop switch. |
| 115 | grad_state = grad_ctxt.grad_state |
| 116 | real_pred = grad_state.history_map.get(pred.name) |
| 117 | if real_pred is None: |
| 118 | # Remember the value of pred for every iteration. |
| 119 | grad_ctxt = grad_state.grad_context |
| 120 | grad_ctxt.Exit() |
| 121 | history_pred = grad_state.AddForwardAccumulator(pred) |
| 122 | grad_ctxt.Enter() |
| 123 | |
| 124 | # Add the stack pop op. If pred.op is in a (outer) CondContext, |
| 125 | # the stack pop will be guarded with a switch. |
| 126 | real_pred = grad_state.AddBackpropAccumulatedValue(history_pred, pred) |
| 127 | grad_state.history_map[pred.name] = real_pred |
| 128 | pred = real_pred |
| 129 | # pylint: disable=protected-access |
| 130 | return control_flow_ops._SwitchRefOrTensor(grad, pred, name="cond_grad") |
| 131 | # pylint: enable=protected-access |
| 132 | else: |
| 133 | num_inputs = len(op.inputs) |
| 134 | cond = [math_ops.equal(op.outputs[1], i) for i in xrange(num_inputs)] |
| 135 | # pylint: disable=protected-access |
| 136 | return [control_flow_ops._SwitchRefOrTensor(grad, cond[i])[1] |
| 137 | for i in xrange(num_inputs)] |
| 138 | # pylint: enable=protected-access |
| 139 | |
| 140 | |
| 141 | @ops.RegisterGradient("RefMerge") |
no test coverage detected