← notesupdated Aug 29, 20262 min read

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.

internalsautodiffjvm
Three parallel arrays — opcode, operand indices, partials — growing left to right.

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.


Discuss: Hacker News · open an issue