def testForwardGradientFunctionUsedByAccumulatorForOps(self):
previous_fn = forwardprop._forward_gradient
try:
with forwardprop.ForwardGradientAccumulator() as acc:
x = constant_op.constant(1.)
acc.watch(x, 2.)
y = x + x
pywrap_tensorflow.TFE_Py_RegisterForwardGradientFunction(
lambda *args, **kwargs: [constant_op.constant(-15.)])
z = x + x
self.assertAllClose(4., acc.jvp(y))
self.assertAllClose(-15., acc.jvp(z))
finally:
pywrap_tensorflow.TFE_Py_RegisterForwardGradientFunction(previous_fn)
transpose_tape.watch(forwardprop_aids)
grads = backfunc_tape.gradient(
trainable_outputs,
trainable_inputs,
forwardprop_aids,
unconnected_gradients=UnconnectedGradients.ZERO)
nontrivial_output_tangents = transpose_tape.gradient(
grads, forwardprop_aids, output_gradients=nontrivial_tangents)
output_tangents = [None] * len(outputs)
for index, tangent in zip(nontrivial_output_indices,
nontrivial_output_tangents):
output_tangents[index] = tangent
return output_tangents
pywrap_tensorflow.TFE_Py_RegisterForwardGradientFunction(_forward_gradient)
class ForwardGradientAccumulator(object):
"""Computes Jacobian-vector products using forward-mode autodiff.
Example:
```
with ForwardGradientAccumulator() as acc:
x = tf.constant([[2.0, 3.0], [1.0, 4.0]])
acc.watch(x, tf.constant([[5., 6.], [7., 8.]]))
y = tf.reduce_sum(tf.sin(x) * tf.tan(x), axis=1)
jvp = acc.jvp(y)
```