A Differentiable Programming System to Bridge ML and Scientific Computing

Observation/request: For higher-order derivatives (Hessian, Laplacian, etc), AD libraries typically provide API shortcuts. I have found it difficult to control or predict the memory footprint and the time complexity of these API shortcuts.

Laplacian is case in point: it is sometimes computed by first computing the Hessian by forward-over-reverse and then taking the trace. In most libraries you would have to dig deep in the documentation or the code itself to understand how intermediary values are accumulated in memory for instance.

I would love to have a summary table of the complexity of all these API shortcuts, similar to https://wiki.python.org/moin/TimeComplexity for data-structures, but for these AD shortcuts (laplacian, divergence, hessian, etc). The table would show memory and time complexity in terms of the number of the input/intermediary/output dimensions and the complexity of evaluating the original function.

Compare that to the situation in Julia: ∂P works today on arbitrary programs—like the ray tracer and other examples in this paper—programs which are highly non-trivial and use iteration, recursion, mutation and global state. All of which Julia's ∂P can take derivatives through.

When you additionally consider Swift's essentially non-existent computational and data science ecosystem, it's a bit hard (for me at least) to rationalize the Swift ∂P effort. (Are we going to differentiate iPhone apps?) They're attempting to bootstrap their computational/data ecosystem by allowing calling Python code, but as soon as you call into Python, you lose all ability to take derivatives which only works for pure Swift code. So any program which relies on Python to do some of the computation you want to take a derivative of won't be differentiable, which kind of defeats the point of having ∂P in the first place. We'll see how it pans out but the Swift effort has considerable technical and social challenges to overcome.

Conceptually, I enjoy the functional approach where the standard compiler does the work. But I'm still using some graph based approach, where the order of operations are recorded on a "tape", then replayed it in reverse order.

The main reason is to be able to manage the memory and computation placement. To distribute huge computations among machines, while limiting memory transfers between machines/devices to a minimum.

Also some kind of automatic gradient accumulation to reduce the memory need would be nice. Currently using the graph based approach, partitioning a big tensor into smaller ones and doing a map-reduce to accumulate the gradient works surprising well (except for the initial very long graph creation time).

Every time I do one of those manual optimizations I start dreaming : "that should be done by the compiler". You tell the compiler your cluster definition and devices, you write a few for loops as if the code was done on a single cpu. And the compiler do the loop splitting, reordering, parallelization across device in an optimal way given your available memory constraints.