Writing a custom op

Every payoff in this Learn section has been built from the same dozen primitives: + - * / exp log sqrt abs max min randn, plus constants. What do you do the day you want a genuinely new building block — a smoothed clamp, a house payoff kink nobody's written yet — without touching the engine at all?

The whole story

The registry is two maps and an interface

Real source, CustomOp.java, in full for the unary half:

@FunctionalInterface
public interface Unary {
  ADouble apply(AadRecorder rec, ADouble x);
}

private static final Map<String, Unary> UNARY = new ConcurrentHashMap<>();

public static void registerUnary(String name, Unary op) {
  UNARY.put(requireName(name), Objects.requireNonNull(op, "op"));
}

public static Unary unary(String name) {
  Unary op = UNARY.get(name);
  if (op == null) {
    throw new IllegalArgumentException("no unary custom op named '" + name + "'; registered: " + UNARY.keySet());
  }
  return op;
}

That's the entire mechanism. registerUnary puts a function in a map; unary(name).apply(rec, x) looks it up and calls it. There's no builder, no annotation processor, no code generation step — a custom op is a name for a lambda, nothing more.

A custom op, end to end

The engine's own cookbook (docs/cookbook/custom-ops.md) registers a soft clamp to [-1, 1] from two calls to Smooth.ramp (Chapter 12.1's rolled kernel used the same Smooth/SpecialFn classes' arithmetic, one chapter earlier — this is what actually sits inside a tape node):

CustomOp.registerUnary("softclamp", (rec, x) ->
    Smooth.ramp(rec, x.add(1.0), 0.05).sub(Smooth.ramp(rec, x.sub(1.0), 0.05)).sub(1.0));

ADouble y = CustomOp.unary("softclamp").apply(rec, someScalar);

Real numbers, this session, one scenario per x, adjoint from the same reverse sweep every other page in this Learn section has used:

x	softclamp(x)	d/dx	region
-2.00	-1.000000	0.000000	saturated low
-1.00	-0.965343	0.500000	at the kink
-0.30	-0.300000	0.999999	linear
0.00	0.000000	1.000000	linear
1.00	0.965343	0.500000	at the kink
2.00	1.000000	0.000000	saturated high

It's exactly x through the linear region, exactly half the neighboring slope at each kink — the two ramps' transitions overlapping — and flat once saturated. Not a plausible shape; the printed table above is System .out from a real run.

It's a macro, not a special op

CustomOp.unary(name).apply(rec, x) is a plain Java method call at record time — it runs the registered lambda once, right there, and whatever primitive nodes that lambda appends are the only trace it leaves. Measured this session, same tape, three ways to build it:

via CustomOp.unary("softclamp"):  nodes=34   (cpu, cpu-jit, simd — all three)
inlined by hand, no registry:     nodes=34

Identical. AadTape has no CUSTOM_OP opcode and never will need one for anything expressible this way — a custom op and the same expression typed out by hand are, to the tape, the same tape. Which is also why it needs nothing from any code generator: cpu-jit's KernelGenerator (Chapter 12.1), simd's vectoriser, and every GPU backend's codegen were already able to differentiate ABS/DIV/EXP/LOG/MUL/MAX before your op existed. Naming a composition of them doesn't ask any of those generators to learn anything new.

Four ops already registered

static {
  registerUnary("relu", (rec, x) -> x.max(0.0));
  registerUnary("softplus", (rec, x) -> Smooth.ramp(rec, x, 0.05));
  registerUnary("sigmoid", (rec, x) -> Smooth.step(rec, x, 1.0));
  registerUnary("normCdf", SpecialFn::normCdf);
}

Worth noticing which one doesn't belong with the other three: relu is the literal, discontinuous max(x, 0) — the same kinked MAX node Chapter 3.2 met on a barrier payoff, not a smoothed stand-in. softplus, sigmoid and normCdf are all genuinely mollified. If your own custom op needs a clean adjoint through the kink, reach for Smooth.ramp the way softplus does, not relu's bare max.

Try it yourself

Change the width in the cookbook's softclamp registration from 0.05 to 0.5 — one number — and rerun. Real numbers, this session: softclamp(1 .0) moves from 0.965343 (barely below the saturated value) to 0.662501, and the derivative at that same kink moves from 0.500000 to 0.482014. A wider smoothing width blurs the clamp's corner further out on both sides — the same width-versus-sharpness trade-off Chapter 3.2's barrier smoothing made, one parameter, immediately visible.

▶️ Run it

The registry, the softclamp registration, and the sweep above — live, available here. CustomOp and Smooth run completely unmodified; this is the one lesson in this whole Learn section whose engine-side code needed no rewrite at all, since there's no reflection anywhere in CustomOp's own mechanism to route around. The one substitution is recording through rec.input() directly instead of the standalone snippet's Nabla.model (new P(x), ...), which still goes through a one-field record's reflective component lookup — kept only for consistency with every earlier chapter's cell, not because this page's own code needs it:

There's no dedicated nablatensor-examples class for a custom op yet — the table above came from a small standalone program built against the engine's own jars, following exactly the pattern OpsTest.java (nablatensor-ops/src/test/java/.../OpsTest.java) uses for its own value-and-adjoint checks:

CustomOp.registerUnary("softclamp", (rec, x) ->
    Smooth.ramp(rec, x.add(1.0), 0.05).sub(Smooth.ramp(rec, x.sub(1.0), 0.05)).sub(1.0));

try (Nabla.TypedPricer<P> pricer = Nabla.model(new P(x), (rec, in) ->
    rec.output(CustomOp.unary("softclamp").apply(rec, in.of(P::x))))
    .fp64().greeks().on("cpu-jit").build()) {
  var v = pricer.value().with(new P(x)).scenarios(1).seed(1L).run();
  System.out.println(v.price() + "  " + v.greek(P::x) + "  nodes=" + pricer.nodes());
}

(P is any one-field record, e.g. record P(double x) {}.) Swap "cpu-jit" for "cpu" or "simd" and the printed value, derivative and node count don't move — that agreement is OpsTest's own customOpRunsIdenticallyOnEveryCpuBackend test, run by hand instead of through Maven.

⚠️ What this doesn't do

OpsTest also has a customOpRunsOnAGpuBackendWithNoEngineEdit test — the Phase-1 definition-of-done the cookbook's "Verification" section describes — but it calls AadEngines.available(...) to find a GPU backend, the exact auto-probe this Learn section has avoided running directly since Chapter 4.1 (see the sandbox-is-real-desktop caution carried since then). Nothing here disputes that the GPU leg passes; it just wasn't re-run on this machine, same precedent as every GPU number in Chapters 4 and this one. This page also stops at the composable "macro" form (Seam 3): a fused op — one that isn't expressible as a composition of existing primitives at all, with a hand-written forward and adjoint rule registered directly into every code generator — is, per CustomOp's own class doc, "a later engine feature," not something this version of the engine can do yet.

What's next

→ Deeper: docs/cookbook/custom-ops.md has the full building-blocks table (Smooth.gt/between, SpecialFn.erf/ pow) this page didn't need for one example. → Next: Anatomy of a GPU kernel source file — a different GPU-facing subsystem, the tensor library's own hand-written CUDA-C catalog.