Mesh & Geometry Invariant Neural Surrogates for Hypersonic Flows

Different Kinds of Aircrafts & Flight Conditions

Predict Quantities of Interest on a 3D Grid

1

FOR MIT & DOD DISTRIBUTION ONLY

Weight-Tied MLPs for Mesh-Invariant PDE Solution

2

FOR MIT & DOD DISTRIBUTION ONLY

Predicted vs Ground Truth Pressure Forces

3

FOR MIT & DOD DISTRIBUTION ONLY

Model Predictions for the Waverider at Unseen Conditions

4

FOR MIT & DOD DISTRIBUTION ONLY

Model Predictions for the Waverider at Unseen Conditions

5

FOR MIT & DOD DISTRIBUTION ONLY

Model Predictions for the Waverider at Unseen Conditions

6

FOR MIT & DOD DISTRIBUTION ONLY

MLPs Generalize to Unseen Flight Conditions

Weight-Tied MLPs Generalize to Unseen Mach Numbers

7

FOR MIT & DOD DISTRIBUTION ONLY

Embedding Continuity Equations

3D Compressible Navier Stokes

Embedding Continuity Equations for 2D Incompressible Navier Stokes

Automatic Differentiation in the Forward Pass makes Training & Inference Slow

8

FOR MIT & DOD DISTRIBUTION ONLY

Standard ML Compilation Pipeline

From the Official StableHLO/OpenXLA Documentation

Google Colab’s Julia Runtime ships with Reactant

& Lux developed as part of this Project.

9

FOR MIT & DOD DISTRIBUTION ONLY

Tensor Optimizations with EnzymeJAX & Reactant

Compile Program to MLIR

Run Optimizations + AutoDiff

Compile & Execute

Code Generation for Heterogeneous Hardware

10

FOR MIT & DOD DISTRIBUTION ONLY

Reactant Compilation Pipeline

Reactant

• Tracing (Abstract Interpreter & Multiple Dispatch)
• Emit MLIR
• Resolve multi-device execution

Enzyme JAX

• Optimization Passes
• Automatic Differentiation
• Automatic Batching (similar to jax.vmap)

11

FOR MIT & DOD DISTRIBUTION ONLY

100+ Additional Compiler Optimizations over XLA

Embedding Continuity Equations for 2D Incompressible Navier Stokes

Compiler Optimizations Cut Down the Total Matrix Multiplications Needed From 12 to 3

4x Improvement in the Core Bottleneck Operations

12

FOR MIT & DOD DISTRIBUTION ONLY

Reactant & EnzymeJAX Accelerate Tensor Workloads

Reactant makes Julia competitive (and even faster than native Jax) for Deep Learning

Even Highly Optimized Workloads like ResNet see a Performance Improvement upon EnzymeJAX compilation

13

FOR MIT & DOD DISTRIBUTION ONLY

Embedding Continuity Equation for 3D Compressible Navier Stokes

Predicting Mass Flux (instead of Velocities) implicitly enforces Continuity Equation

Forward Pass requires 6 partial derivatives. Scaling is impossible without Reactant!

14

FOR MIT & DOD DISTRIBUTION ONLY

Preliminary Results Validate Accurate Predictions

Embedded Continuity Equations enforce Physical Constraints and Generalize to Unseen Flight Conditions

15

FOR MIT & DOD DISTRIBUTION ONLY

Compiler Optimizations Enable Fast Predictions

CPU

GPU

Weight-Tied MLP

Continuity Conserving MLP

Predictions on a 1024x1024 Grid

~7-8s

~20-21s

~0.047s

~0.127s

16

FOR MIT & DOD DISTRIBUTION ONLY

Ongoing Work: Geometry Invariant Predictions

Use Surface Mesh Embeddings to Condition the Predictions

17

FOR MIT & DOD DISTRIBUTION ONLY