Recording the tape as a flat array
The data structure under adjoint AD is usually a graph of heap nodes. We use a struct-of-arrays trace instead, and the JVM rewards it with predictable cache behaviour.
Most adjoint AD libraries model the tape as a linked structure: every intermediate value is an object that points at its parents and carries its local partial derivatives. It is easy to write and it allocates one small object per operation. For a Monte-Carlo book that is tens of millions of operations per path, and the garbage collector notices.
Struct of arrays
NablaTensor records the tape as a handful of primitive arrays that grow together:
final class Tape {
int[] op; // opcode: ADD, MUL, EXP, LOG, ...
int[] lhs, rhs; // indices of operands earlier in the tape
double[] value; // forward value at this slot
double[] adjoint; // reverse accumulator, zeroed before the sweep
int length;
int record(int opcode, int a, int b, double v) {
op[length] = opcode; lhs[length] = a; rhs[length] = b;
value[length] = v;
return length++;
}
}
A slot is 40 bytes and, crucially, the arrays are contiguous. The forward sweep
writes value front to back; the reverse sweep reads op/lhs/rhs back to
front and scatters into adjoint. Both passes are linear scans, which is the
access pattern hardware prefetchers are built for.
The reverse sweep
void reverse() {
java.util.Arrays.fill(adjoint, 0, length, 0.0);
adjoint[length - 1] = 1.0; // seed: d(result)/d(result)
for (int i = length - 1; i >= 0; i--) {
double a = adjoint[i];
if (a == 0.0) continue;
switch (op[i]) {
case ADD -> { adjoint[lhs[i]] += a; adjoint[rhs[i]] += a; }
case MUL -> {
adjoint[lhs[i]] += a * value[rhs[i]];
adjoint[rhs[i]] += a * value[lhs[i]];
}
case EXP -> adjoint[lhs[i]] += a * value[i];
// ...
}
}
}
No local-partial storage: for cheap primitives it is faster to recompute the
partial from the recorded forward value than to have written it down.
The switch is the hot loop. Keeping the opcode set small (about a dozen primitives) is what lets the JIT turn it into a jump table instead of a chain of comparisons. Every convenience primitive you add has a cost here.
Next note: replaying one tape across a million scenarios without rebuilding it.