Gradients for a Switch op is calculated using a Merge op. If the switch is a loop switch, it will be visited twice. We create the merge on the first visit, and update the other input of the merge on the second visit. A next_iteration is also added on second visit.
(op, *grad)
| 33 | |
| 34 | |
| 35 | def _SwitchGrad(op, *grad): |
| 36 | """Gradients for a Switch op is calculated using a Merge op. |
| 37 | |
| 38 | If the switch is a loop switch, it will be visited twice. We create |
| 39 | the merge on the first visit, and update the other input of the merge |
| 40 | on the second visit. A next_iteration is also added on second visit. |
| 41 | """ |
| 42 | graph = ops.get_default_graph() |
| 43 | # pylint: disable=protected-access |
| 44 | op_ctxt = op._get_control_flow_context() |
| 45 | grad_ctxt = graph._get_control_flow_context() |
| 46 | # pylint: enable=protected-access |
| 47 | if isinstance(op_ctxt, WhileContext): |
| 48 | merge_grad = grad_ctxt.grad_state.switch_map.get(op) |
| 49 | if merge_grad is not None: |
| 50 | # This is the second time this Switch is visited. It comes from |
| 51 | # the non-exit branch of the Switch, so update the second input |
| 52 | # to the Merge. |
| 53 | # TODO(yuanbyu): Perform shape inference with this new input. |
| 54 | if grad[1] is not None: |
| 55 | # pylint: disable=protected-access |
| 56 | control_flow_ops._AddNextAndBackEdge(merge_grad, grad[1], |
| 57 | enforce_shape_invariant=False) |
| 58 | # pylint: enable=protected-access |
| 59 | return None, None |
| 60 | elif grad[0] is not None: |
| 61 | # This is the first time this Switch is visited. It comes from |
| 62 | # the Exit branch, which is grad[0]. grad[1] is empty at this point. |
| 63 | # Use grad[0] for both inputs to merge for now, but update the second |
| 64 | # input of merge when we see this Switch the second time. |
| 65 | merge_grad = merge([grad[0], grad[0]], name="b_switch")[0] |
| 66 | grad_ctxt.grad_state.switch_map[op] = merge_grad |
| 67 | return merge_grad, None |
| 68 | else: |
| 69 | # This is the first time this Switch is visited. It comes from the |
| 70 | # Identity branch. Such a Switch has `None` gradient for the Exit branch, |
| 71 | # meaning the output is not differentiable. |
| 72 | return None, None |
| 73 | elif isinstance(op_ctxt, CondContext): |
| 74 | zero_grad = grad[1 - op_ctxt.branch] |
| 75 | # At this point, we have created zero_grad guarded by the right switch. |
| 76 | # Unfortunately, we may still get None here for not trainable data types. |
| 77 | if zero_grad is None: |
| 78 | # For resource variables we get None always on the other branch, so bypass |
| 79 | # this. |
| 80 | if op.inputs[0].dtype == dtypes.resource: |
| 81 | return merge( |
| 82 | [grad[op_ctxt.branch]] * 2, name="cond_resource_grad")[0], None |
| 83 | return None, None |
| 84 | return merge(grad, name="cond_grad")[0], None |
| 85 | else: |
| 86 | false_grad = switch(grad[0], op.inputs[1])[0] |
| 87 | true_grad = switch(grad[1], op.inputs[1])[1] |
| 88 | return merge([false_grad, true_grad])[0], None |
| 89 | |
| 90 | |
| 91 | ops.RegisterGradient("Switch")(_SwitchGrad) |
nothing calls this directly
no test coverage detected