add pedagogically sound examples for custom derivatives #1473
Conversation
| @@ -51,7 +51,7 @@ double invmag(double* __restrict__ A, int n) { | |||
| return Q_rsqrt(sumsq); | |||
| } | |||
|
|
|||
|
|
|||
| //// (1) An interface that just happens to work because Q_rsqrt is simple enough. | |||
There was a problem hiding this comment.
So the issue here isn't because of Q_sqrt, but rather it is because of the use of Q_sqrt within the function being differentated.
In particular, someones one may need the original primal return to compute the derivative of a function. e.g. if we did expr(Q_sqrt(x)). However, since we simply return Q_rsqrt (and the default semantics of autodiff without a flag is to just return the derivative, not the primal), we didn't need it.
In the second case, becuase it is stored into a duplicated, the same value must be stored in the derivative (and thus the primal return is required). If the variable A was marked dup_noneed, then the return-less version would've been fine.
The julia custom rule API goes into more details: https://enzyme.mit.edu/index.fcgi/julia/stable/generated/custom_rule/
That said this is very undocumented, and counter-intuitive.
@samuelpmish and I are working on trying to make a nicer c/c++ interface, including for custom derivatives so this is a great example to have.
There was a problem hiding this comment.
Ah, yes. I see. It is nice to know what legalCombinedForwardReverse does now.
For some reason, replacing the return with return __builtin_exp(Q_sqrt(sumsq)) wasn't enough to trigger the assert, but additionally deleting the definition of Q_sqrt and replacing it with extern float Q_sqrt(float); was.
| double x0 = *(double*)tape; // x points to junk; original input *x must be obtained from the tape | ||
| *x_b -= (*y_b) * Q_rsqrt(x0) / (2 * x0); | ||
| *y_b = 0; // Since y is used purely as an output, changes in y do not affect the calculation | ||
| free(tape); |
There was a problem hiding this comment.
This is perhaps off-topic, but what if the user wishes to apply reverse mode multiple times? Should freeing the tape not be decoupled from the numerics of the reverse pass?
There was a problem hiding this comment.
If the code is generated by Enzyme, the no-free flag will keep that cache around. We need to figure out a nice interface for specifying that (and other options) for the custom derivatives.
Alternatively, perhaps we just have a special allocator that manages this behind the scenes.
There was a problem hiding this comment.
Also at least in this case, one doesn't need to return a void*, one can just use a double, rather than double* as the tape.
There was a problem hiding this comment.
I agree that in this case the tape can be of type double, but a void star will get you a lot further in more general cases. I will respond to the initial question. It is my understanding that when one writes an augmented function (f_fwd) and a reverse function (f_rev), one is writing a pair of coupled functions; every call to f_fsd is expected to eventually followed up by precisely one corresponding call to f_rev - and with the corresponding tapes. Given this interface, there this no other place but f_rev to cleanup whatever mess was made by f_fwd. Furthermore, since the same code is run for multiple calls to f (suppose you want to differentiate a function g which calls f twice), the tape is the only way in general to distinguish the two separate pieces of data that f_rev will need.
I am not sure what this no-free flag is. However, I do know the enzyme_allocated option can handle fairly general cases - but this option is not (and cannot be?) integrated with the custom derivatives.
resolves #1395
As the integration folder seems to be the best documentation of the C/C++ interface, it is important to have high-quality examples here.