[improbable22]

This looks interesting, thanks.

What you call "delimited continuations" sounds a bit like how AD works in Julia's Flux package (and maybe elsewhere): During the forward calculation, a chain of functions is constructed, whose evaluation is the backward pass. This is done by overloading the original * to return both x * y and a closure Δ -> (Δ * y, x * Δ).

Does that sound right, are these indeed similar, or have I mis-understood something? I have never read Scala, but if I squint I can make your overloading of * is doing something similar.

[tiarkrompf]

It's different -- instead of returning a closure, we take a closure as additional parameter (check paper for details). This means that the call stack stays intact and all intermediate values can be stack-allocated.

[dzheng]

Could you please elaborate on what exactly enables stack allocation of intermediate values?

Is it a direct result of the semantics of delimited continuations? (perhaps because closures are passed as parameters and don't capture variables)

Or is it enabled via staging or optimizations?

[one-more-minute]

It's specific to CPS form: every basic statement is a full function call, which moves to the next statement by calling another function (i.e. the continuation). Normally this would just be a very quick way to get yourself a stack overflow, so typical CPS-form compilers optimise (or require) the tail-call case, where you can pop the stack frame for the last instruction before moving on to the next. But for AD that's not a big deal, you just re-use those values in the backwards pass.

I doubt it's better than just allocating your own stack on the heap, especially when you have something like a differential equation solver with millions of statements, but it's still a neat trick.

[dzheng]

Aha, that makes sense, thanks!

I agree that the performance benefits of such stack allocation (over heap allocation) aren't quite clear in practice.

I feel the bigger win of delimited cont./closure-based AD approaches is that they can model the control flow of reverse-mode AD without AD-specific code transformations. Delimited cont. is especially great at making things modular: each differentiable function performs primal computation, calls the callback with primal result, then performs adjoint computation.

[improbable22]

OK, thanks, sounds worth trying to wrap my head around...

(Note BTW that the github page has a dead link to the paper.)

[tiarkrompf]

Thanks! Fixed now.