The python grad function for the Forward function.
(op, *args)
| 377 | [theta, state1, inputs, acc_state, acc_extras]) |
| 378 | |
| 379 | def Grad(op, *args): |
| 380 | """The python grad function for the Forward function.""" |
| 381 | |
| 382 | # NOTE: tf.gradient backprops None for int32/int64 while zeros |
| 383 | # for float32/float64. For consistency, we always backprop |
| 384 | # zeros. |
| 385 | args = list(args) |
| 386 | for i, dy in enumerate(args): |
| 387 | if dy is None: |
| 388 | args[i] = array_ops.zeros_like(op.outputs[i]) |
| 389 | # TODO(drpng): getting the extra state here? |
| 390 | op_inputs = [x for x in op.inputs] |
| 391 | op_struct = [ |
| 392 | self._theta, self._state, self._inputs, self._max_input_length, |
| 393 | self._extras |
| 394 | ] |
| 395 | (theta, state0, inputs, max_input_length, _) = _Pack(op_inputs, op_struct) |
| 396 | # acc_state and acc_extras are computed by the Forward pass and |
| 397 | # needed by the Backward pass. |
| 398 | acc_state, _, acc_extras = _Pack([x for x in op.outputs], |
| 399 | [self._state, self._state, self._extras]) |
| 400 | |
| 401 | # Forward computes acc_state, the final state and |
| 402 | # acc_extras. tf.gradients gives us their gradients w.r.t. the |
| 403 | # final loss. Because acc_extras are not exposed by Compute(), |
| 404 | # it has no gradients w.r.t. the final loss (i.e., by |
| 405 | # construction, it must be zeros). |
| 406 | d_acc_state, d_state1, _ = _Pack(args, |
| 407 | [self._state, self._state, self._extras]) |
| 408 | return Backward(*_Flatten([ |
| 409 | theta, state0, inputs, max_input_length, acc_state, acc_extras, |
| 410 | d_acc_state, d_state1 |
| 411 | ])) |
| 412 | |
| 413 | # Forward calls ForwardLoopBody n times. Each time computes one |
| 414 | # time step of the recurrent net. |