A new approach allows for backpropagation through existing Fortran, C, or C++ simulation code, embedding it into JAX and PyTorch as a high-performance differentiable physics engine. This bridges the gap for decades of validated physics code in fields like CFD and climate modeling, which typically lack gradient information.

Instead of rewriting code, the method leverages Enzyme, an automatic differentiation tool that operates at the LLVM IR level. By combining LFortran (a Fortran compiler), LLVM, and Enzyme, exact gradients can be obtained from Fortran code. The process involves compiling the Fortran solver to LLVM IR, applying Enzyme for differentiation, and then wrapping the result as a custom JAX primitive. This experimental approach, though requiring effort to debug, enables gradients through complex, multi-step simulations, matching analytic answers.

The system was tested on a 220-line Fortran 90 heat solver. Challenges included avoiding Fortran features that compile to runtime calls, which Enzyme cannot trace. The compilation pipeline involves six LLVM optimization passes, with key steps being the LFortran IR generation, a pre-Enzyme optimization pass, and the Enzyme differentiation pass itself. LFortran was chosen over Flang due to its cleaner LLVM IR output, which is more amenable to Enzyme's analysis.