Core mechanism
The primitive supplies its numerical evaluation and the contribution it makes to each input adjoint. The recorder can then compose it with existing nodes and every backend can replay the agreed semantics.
Engine Internals & GPU Performance
Extend the Tape With a New Differentiable Primitive
A registered custom op expands into ordinary tape primitives during recording. Java registers swish at runtime, then compares every unary op's recorded value and adjoint with a central bump.
Evaluate at
Swish is registered as x × sigmoid(x); no engine-specific derivative is supplied.
This exact source runs in TeaVM. Form changes update its Java literals and reset manual edits.
import com.nablatensor.engine.ADouble;
import com.nablatensor.engine.AadEngines;
import com.nablatensor.engine.AadExecutable;
import com.nablatensor.engine.AadOptions;
import com.nablatensor.engine.AadRecorder;
import com.nablatensor.engine.AadResult;
import com.nablatensor.engine.AadTape;
import com.nablatensor.ops.CustomOp;
public final class CustomOpRiskStudio {
private CustomOpRiskStudio() {}
public static void main(String[] args) {
double x0 = 0.3, bumpH = 0.001;
CustomOp.registerUnary("swish", (rec, x) -> x.mul(CustomOp.unary("sigmoid")
.apply(rec, x)));
String[] names = {
"relu", "softplus", "sigmoid", "normCdf", "swish"
};
AadOptions options = AadOptions.defaults()
.withPrecision(AadOptions.PrecisionEnum.FLOAT64)
.withThreads(1);
for (String name : names) {
AadTape tape = AadRecorder.record(rec -> {
ADouble x = rec.input("x", x0); rec.output(CustomOp.unary(name)
.apply(rec, x));
});
try (AadExecutable exec = AadEngines.require("cpu", options)
.compile(tape, options)) {
exec.setInput("x", x0);
AadResult base = exec.replaySafe(1L, 42L);
exec.setInput("x", x0 + bumpH);
double up = exec.replaySafe(1L, 42L)
.value();
exec.setInput("x", x0 - bumpH);
double down = exec.replaySafe(1L, 42L)
.value();
double bump = (up - down) / (2.0 * bumpH);
System.out.println("ROW|" + name + "|" + tape.size() + "|" + base.value() + "|" + base.gradient("x")
+ "|" + bump);
}
}
}
}
TeaVM compiles and runs the Java source above in this browser.
Implementation guide
Extending the graph with a differentiable primitive
A custom operation becomes safe only when its forward value and reverse derivative rule are treated as one contract.
Practical workflow
Derive the adjoint algebraically, test values and derivatives against finite differences, define domain/error behaviour, and add backend conformance tests before using the op in products.
Key details
*Keywords: custom operation automatic differentiation java, extend aad engine, register op monte carlo, differentiable macro*
The engine's primitive op set is + - * / neg exp log sqrt abs max min plus randn and constants. Everything a quant payoff needs is a composition of those, so a "custom op" in NablaTensor is a macro: a function that expands into primitive nodes when it is recorded. Its adjoint is whatever the recorded sub-graph produces — no special reverse rule, and it runs unchanged on cpu-jit, simd and every GPU backend.
OpsTest records each op as a one-scenario deterministic function and checks the value and the adjoint against the closed form, checks that shrinking a smoothing width recovers the discontinuous limit, and runs a user-registered op identically on cpu, cpu-jit, simd and a GPU backend — the Phase-1 definition-of-done for custom ops.
A genuinely non-composable kernel — one that must call an external special function not expressible in the primitive set — would need a new AadOpEnum value handled in every code generator (scalar interpreter, batched, bytecode, SIMD, CUDA-C, GLSL). That fused {forward, adjoint} form is a planned engine feature; in practice the special functions quants reach for (N(x), erf, pow, smoothed indicators, softplus, …) are all composable and already here.
Scope and review point
An incorrect adjoint can produce plausible prices and wrong risk. Numerical tests and independent review are mandatory for each new primitive.