Gradients for an Enter are calculated using an Exit op. For loop variables, grad is the gradient so just add an exit. For loop invariants, we need to add an accumulator loop.
(op, grad)
| 201 | |
| 202 | @ops.RegisterGradient("Enter") |
| 203 | def _EnterGrad(op, grad): |
| 204 | """Gradients for an Enter are calculated using an Exit op. |
| 205 | |
| 206 | For loop variables, grad is the gradient so just add an exit. |
| 207 | For loop invariants, we need to add an accumulator loop. |
| 208 | """ |
| 209 | graph = ops.get_default_graph() |
| 210 | # pylint: disable=protected-access |
| 211 | grad_ctxt = graph._get_control_flow_context() |
| 212 | # pylint: enable=protected-access |
| 213 | if not grad_ctxt.back_prop: |
| 214 | # Skip gradient computation, if the attribute `back_prop` is false. |
| 215 | return grad |
| 216 | if grad_ctxt.grad_state is None: |
| 217 | # Pass the gradient through if we are not in a gradient while context. |
| 218 | return grad |
| 219 | if op.get_attr("is_constant"): |
| 220 | # Add a gradient accumulator for each loop invariant. |
| 221 | if isinstance(grad, ops.Tensor): |
| 222 | result = grad_ctxt.AddBackpropAccumulator(op, grad) |
| 223 | elif isinstance(grad, ops.IndexedSlices): |
| 224 | result = grad_ctxt.AddBackpropIndexedSlicesAccumulator(op, grad) |
| 225 | else: |
| 226 | # TODO(yuanbyu, lukasr): Add support for SparseTensor. |
| 227 | raise TypeError("Type %s not supported" % type(grad)) |
| 228 | else: |
| 229 | result = exit(grad) |
| 230 | grad_ctxt.loop_exits.append(result) |
| 231 | grad_ctxt.ExitResult([result]) |
| 232 | return result |
| 233 | |
| 234 | |
| 235 | @ops.RegisterGradient("RefEnter") |
no test coverage detected