Enzyme
Enzyme copied to clipboard
Add Jacobian mode
Add a derivative mode __enzyme_jacobian
, that computes the jacobian using forward and reverse mode.
@tgymnich Could you give a quick summary of this issue ? (viz. Requirements)
The Jacobian mode can be seen as an addition to the LLVM/Clang frontend for Enzyme, which is making use of the scalar reverse or forward modes and the vector reverse or forward modes to compute the vector Jacobian product / Jacobian vector product respectively. There already exists an implementation for this in Enzyme.jl you can refer to.
The whole thing should roughly work like this in the end:
constexpr int N = 2; // num columns of J
constexpr int M = 3; // num rows of J
extern void __enzmye_fwd_jacobian(...);
extern void __enzmye_rev_jacobian(...);
// Variables serving as metadata to mark the argument that follows
extern int batch_size;
extern int num_outs;
void f(double *v, double* res) {
res[0] = v[1];
res[1] = v[0] * v[0];
res[2] = v[0] * v[0] * v[0]
}
double x[N] = [2.0, 3.0];
double shadow[N][N] = [[1.0, 0.0], [0.0, 1.0]]; // onehot encoded version of x
double JvP[M][N] = {0};
double vJP[M][N] = {0};
// shadow should be an optional argument which defaults to onehot(x).
// batch_size and num_outs should be compile time constants.
__enzmye_fwd_jacobian(f, enzyme_dup &x, &shadow, &JvP, batch_size, N);
__enzmye_rev_jacobian(f, &x, &vJP, num_outs, Mm batch_size, N);
double [M][N] expected_result = [[0.0, 1.0], [4.0, 0.0], [12.0, 0.0]];