JAX-Fluids: Toward Differentiable CFD of Compressible Single- and Two-phase Flows

JIT compilation in JAX comes with specific requirements. One crucial requirement is that array shapes must be fixed at compile time. Once a function is compiled using JIT, the arrays become traced, meaning that JAX recognizes only the shape and data type of the arrays, not their actual values. To implement control flow statements that depend on the actual values of input arguments, you must use the static arguments keyword.

In our example, we illustrate a JIT-compiled function that takes a 3D buffer as input, along with cell sizes and an additional argument specifying the axis direction for derivative computation. The number of halo cells must also be provided. The axis direction must be specified as a static argument, as JAX requires this information at compile time to implement control flow based on its value. Similarly, the number of halo cells must be marked as static so that JAX can infer the shapes of all generated buffers at compile time.

These restrictions necessitate special considerations for implementing spatial conditions in compiled functions, particularly through the use of masks. For instance, in an upwind finite difference implementation, we compute first derivatives based on the local velocity field. The function receives a 1D buffer, the velocity, cell sizes, and the number of cells. We create slice objects for both left and right-sided derivatives and define a mask based on the velocity, where the mask is set to one in regions where the velocity is greater than zero and zero elsewhere. We then use JNP.wa to combine the two derivative options, ensuring that shapes are fixed and known at compile time.

Next, we explore useful function transformations in JAX, starting with JIT compilation. JAX also provides JAX.vmap, which vectorizes functions across a single device. This allows us to evaluate functions on batches of data efficiently. For example, the linear interpolation function takes a vector of points and a computational grid buffer as input. While it can be evaluated for a single point, it is more efficient to evaluate it for millions of points in parallel. Instead of looping through each point sequentially, which would be inefficient on a GPU, we can utilize vmap to specify the axis for broadcasting across all input arguments.

Additionally, we have JAX.pmap, which enables parallel evaluation of functions across multiple devices. For example, if we have a large array with shape N x N, we can split this array according to the number of available devices. By specifying the axis for broadcasting, JAX will distribute the array across the devices and evaluate it in parallel.

Finally, we will also discuss the computation of gradients, which is another important function transformation in JAX.