Re:Foundations for Efficient and Expressive Differentiable P...
by dzheng on Thursday December 27, 2018 @12:43AM
(#18767615)
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.
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.