Gradients for an exit op are calculated using an Enter op.
(op, grad)
| 145 | |
| 146 | @ops.RegisterGradient("Exit") |
| 147 | def _ExitGrad(op, grad): |
| 148 | """Gradients for an exit op are calculated using an Enter op.""" |
| 149 | graph = ops.get_default_graph() |
| 150 | # pylint: disable=protected-access |
| 151 | op_ctxt = op._get_control_flow_context() |
| 152 | grad_ctxt = graph._get_control_flow_context() |
| 153 | # pylint: enable=protected-access |
| 154 | if not grad_ctxt.back_prop: |
| 155 | # The flag `back_prop` is set by users to suppress gradient |
| 156 | # computation for this loop. If the attribute `back_prop` is false, |
| 157 | # no gradient computation. |
| 158 | return None |
| 159 | |
| 160 | if op_ctxt.grad_state: |
| 161 | raise TypeError("Second-order gradient for while loops not supported.") |
| 162 | |
| 163 | if isinstance(grad, ops.Tensor): |
| 164 | grad_ctxt.AddName(grad.name) |
| 165 | else: |
| 166 | if not isinstance(grad, (ops.IndexedSlices, sparse_tensor.SparseTensor)): |
| 167 | raise TypeError("Type %s not supported" % type(grad)) |
| 168 | grad_ctxt.AddName(grad.values.name) |
| 169 | grad_ctxt.AddName(grad.indices.name) |
| 170 | dense_shape = grad.dense_shape |
| 171 | if dense_shape is not None: |
| 172 | grad_ctxt.AddName(dense_shape.name) |
| 173 | grad_ctxt.Enter() |
| 174 | # pylint: disable=protected-access |
| 175 | result = control_flow_ops._Enter( |
| 176 | grad, grad_ctxt.name, is_constant=False, |
| 177 | parallel_iterations=grad_ctxt.parallel_iterations, |
| 178 | name="b_exit") |
| 179 | # pylint: enable=protected-access |
| 180 | grad_ctxt.loop_enters.append(result) |
| 181 | grad_ctxt.Exit() |
| 182 | return result |
| 183 | |
| 184 | |
| 185 | ops.RegisterGradient("RefExit")(_ExitGrad) |