Unlocking Differentiable Fortran: The LFortran and Enzyme Synergy

Posted on 2026-07-09 by @dionhaefner

Imagine a world where you could perform backpropagation directly through legacy Fortran, C, or C++ simulation engines, integrating them into JAX or PyTorch as high-performance, differentiable physics layers.

It is actually possible—provided you have the stomach for experimental tooling. For decades, the gold standard of validated physics code in fields like nuclear engineering, aerospace, and CFD has been locked away in languages that don't expose gradients, creating a massive barrier for modern machine learning pipelines.

Instead of rewriting these systems, we can extract exact gradients using "LLVM-level magic."

The Secret Sauce: Enzyme and LLVM

The breakthrough here is Enzyme, which performs automatic differentiation (AD) at the LLVM Intermediate Representation (IR) level. Because Enzyme operates on the IR, it can differentiate any code that compiles to LLVM.

By "duct-taping" LFortran, LLVM, and Enzyme together, we can process a Fortran thermal solver and produce exact gradients. Tesseract then wraps this output as a custom JAX primitive, effectively turning a legacy Fortran subroutine into a differentiable layer within a JAX graph.

⚠️ Warning: This workflow is highly experimental. Expect to spend significant time debugging NaN gradients and manually inspecting LLVM IR differences to ensure correctness.

However, the payoff is immense: achieving gradients through complex, multi-step time loops that match analytic solutions. It is a fascinating marriage of the oldest and newest technologies in computing.

---

The Struggle with Legacy Gradients

In scientific computing, the need for gradients (derivatives of outputs relative to inputs) is common for optimization and inverse problems. Traditionally, you had three bad options:

The alternative? Automatic compilation of derivatives from the original source.

---

Case Study: A Fortran Heat Solver

To test this, we used thermal_2d.f90, a ~220-line vanilla Fortran 90 program. It simulates 2D transient heat conduction where conductivity depends on temperature.

The Physics

The governing equation is: $\rho \, c_p \frac{\partial T}{\partial t} = \nabla \cdot \big( k(T) \, \nabla T \big) + Q$

Where the material model for conductivity $k$ is linear: $k(T) = k_0 + k_1 \cdot T$

The system uses an explicit Euler method for time integration over n_steps.

The Implementation

The core logic resides in the thermal_2d_solve subroutine. Below is a simplified look at the signature and the internal stencil loop:

subroutine thermal_2d_solve ( n , nx , ny , n_steps , T_init , T_final , T_cur , T_new , k0 , k1 , rho , cp , h_conv , T_inf , T_hot , Q , Lx , Ly , dt )
    implicit none
    integer , intent ( in ) :: n , nx , ny , n_steps
    double precision , intent ( in ) :: T_init ( n )
    double precision , intent ( out ) :: T_final ( n )

    ! Time integration loop
    do step = 1 , n_steps
        do j = 2 , ny - 1
            do i = 2 , nx - 1
                idx = ( j - 1 ) * nx + i
                T_c = T_cur ( idx )
                T_e = T_cur ( idx + 1 )
                T_w = T_cur ( idx - 1 )
                T_nn = T_cur ( idx + nx )
                T_s = T_cur ( idx - nx )

                ! Harmonic-mean conductivity at cell faces
                kx_east = 2.0d0 * ( k0 + k1 * T_c ) * ( k0 + k1 * T_e ) / (( k0 + k1 * T_c ) + ( k0 + k1 * T_e ))
                kx_west = 2.0d0 * ( k0 + k1 * T_c ) * ( k0 + k1 * T_w ) / (( k0 + k1 * T_c ) + ( k0 + k1 * T_w ))

                ! (ky_north, ky_south calculated similarly)
                flux_x = ( kx_east * ( T_e - T_c ) - kx_west * ( T_c - T_w )) / ( dx * dx )
                flux_y = ( ky_north * ( T_nn - T_c ) - ky_south * ( T_c - T_s )) / ( dy * dy )
                T_new ( idx ) = T_c + dt / ( rho * cp ) * ( flux_x + flux_y + Q )
            end do
        end do
    end do
end subroutine

The Compilation Pipeline

The process of making this code differentiable follows this flow:

Project Checklist

• Write vanilla Fortran 90 solver
• Compile to LLVM IR via LFortran
• Apply Enzyme autodiff
• Wrap as JAX primitive via Tesseract
• Solve complex inverse problem (Next Step)