Differentiable systems over framework boundaries: AD interoperability in practice

Speaker: Jonathan Brodrick, Pasteur Labs

Abstract

The diverse range of potential applications offered by automatic differentiation is being unlocked and exploited in many different communities each with their own focus and preferred framework (JAX, Julia, PyTorch, etc.). Established simulation codes increasingly offer derivatives, whether through AD, hand- written adjoints or sensitivity analysis but lack a way to easily expose and compose them with other differentiable components. We believe this has discouraged the research community from fully leveraging available infrastructure; end-to-end gradients over an entire diverse pipeline, from mesher through coupled solvers to surrogates and post-processing, are simply not on the menu. We ask what happens when building such “differentiable systems” becomes practically feasible. This is made tractable by Tesseract, an open source ecosystem that packages any differentiable component behind a simple contract. Tesseracts can then be served and exposed via network through containerisation or independent subprocesses. Reusable recipes enable straightforward encapsulation of JAX and PyTorch workflows and further worked examples show how Julia and Enzyme-equipped Fortran are packaged into Tesseracts, with minimal communication overhead. Specifically, Tesseract-JAX enables forward and reverse-mode AD, batching and partial differentiation over Tesseracts and arbitrary JAX code.

To show that our simple contract supports nontrivial applications, we take Devito, a symbolic finite-difference framework used in industrial seismic imaging, package its adjoint-equipped wave propagators as a Tesseract, and expose them to JAX unchanged. Specifically, we perform seismic full-waveform inversion with simultaneous source estimation, recovering a wave-speed field from recorded traces and unknown source signatures. The wave equation is linear in the source, so the separable sources can be eliminated analytically by forthcoming variable projection functionality in the JAX package lineax, with an outer L-BFGS loop enabling rapid refinement of data misfit from a sensible starting point.

We are also wishing to learn what the EuroAD community would be excited to build if you could differentiate everything and what additional tooling would further enable this.