Enzyme icon indicating copy to clipboard operation
Enzyme copied to clipboard

Add Jacobian mode

Open tgymnich opened this issue 2 years ago • 2 comments

Add a derivative mode __enzyme_jacobian, that computes the jacobian using forward and reverse mode.

tgymnich avatar Mar 14 '22 16:03 tgymnich

@tgymnich Could you give a quick summary of this issue ? (viz. Requirements)

dynamic-queries avatar Nov 27 '22 09:11 dynamic-queries

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]];

tgymnich avatar Nov 28 '22 20:11 tgymnich