def testRecordingSelectively(self): c = constant_op.constant(1.) c_tangent = constant_op.constant(2.) with forwardprop.ForwardAccumulator(c, c_tangent) as acc: with backprop.GradientTape(persistent=True) as tape: tape.watch(c) with tape_lib.stop_recording(): two = constant_op.constant(2.) d = c * two three = constant_op.constant(3.) e = c * three self.assertIsNone(acc.jvp(d)) self.assertIsNone(acc.jvp(e)) self.assertIsNone(tape.gradient(d, c)) self.assertIsNone(tape.gradient(e, c)) tape_lib.record_operation_forwardprop_only( "CustomForwardMul", [d], [c, two], lambda dd: (two * dd, c * dd), None) tape_lib.record_operation_backprop_only( "CustomBackwardMul", [e], [c, three], lambda de: (three * de, c * de)) self.assertAllClose(4., acc.jvp(d)) self.assertIsNone(acc.jvp(e)) self.assertIsNone(tape.gradient(d, c)) self.assertAllClose(3., tape.gradient(e, c))
def testRecordingWithJVPIndices(self): c = constant_op.constant(1.) with forwardprop.ForwardAccumulator(c, 10.) as acc: packed_input_tangents = forwardprop_util.pack_tangents([c ]).tangents self.assertAllClose([10.], packed_input_tangents) d = constant_op.constant(2.) d_tangent = constant_op.constant(3.) tape_lib.record_operation_forwardprop_only( "FunctionWithInlineJVPs", [d] + [d_tangent], [c] + packed_input_tangents, None, (((0, 1), ), )) self.assertAllClose(3., acc.jvp(d))
def testRecordingWithJVPIndices(self): with forwardprop.ForwardGradientAccumulator() as acc: c = constant_op.constant(1.) acc.watch(c, 10.) _, packed_input_tangents = ( pywrap_tensorflow.TFE_Py_PackForwardGradients([c])) self.assertAllClose([10.], packed_input_tangents) d = constant_op.constant(2.) d_tangent = constant_op.constant(3.) tape_lib.record_operation_forwardprop_only( "FunctionWithInlineJVPs", [d] + [d_tangent], [c] + packed_input_tangents, None, (((0, 1), ), )) self.assertAllClose(3., acc.jvp(d))