Mosaic: A Benchmark Suite for Differentiable Physics Solvers

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

Mosaic: A Benchmark Suite for Differentiable Physics Solvers

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

05-2026

arXiv:2606.27895v1 [physics.comp-ph] 26 Jun 2026

Mosaic: A Benchmark Suite for Differentiable Physics Solvers Andrin Rehmann , Heiko Zimmermann , Dion Häfner Pasteur Labs, Brooklyn, NY, USA Correspondence: [email protected]

Abstract Differentiable partial differential equation (PDE) solvers underpin solver-in-the-loop ML training, gradient-based optimal control, and inverse problems, yet the practical cost of obtaining correct, usable gradients from a given solver on a given problem is largely undocumented. Integration effort, computational cost, gradient accuracy, and numerical conditioning vary widely across solvers and are discoverable only by trial and error. We introduce Mosaic, an extensible benchmarking framework for differentiable PDE solvers that standardizes access to solver gradients. Each solver is packaged as a containerized component (Tesseract) exposing a uniform gradient API regardless of language or automatic differentiation (AD) strategy, enabling researchers to evaluate, compare, and build on non-trivial physical solvers. Our evaluation of 14 solvers across fluid dynamics, structural mechanics, and heat transfer demonstrates that the benchmark surfaces practically relevant differences: order-of-magnitude variation in computational cost and Jacobian conditioning, alongside structural incompatibilities that eliminate solvers from realistic tasks entirely. Despite this variation, all solvers that produce gradients converge to similar optima, indicating that the practical barriers are memory limits, numerical stability, and setup compatibility rather than gradient accuracy alone. Mosaic is open-source and available at https://github.com/pasteurlabs/mosaic.

Solver backends

Standardized interface

Benchmark tasks

Evaluation results

forward solve Fwd

PhiFlow

Opt.

XLB

PICT

INS.jl

Grad

PhiFlow

XLB

FEniCS

VJP

Tesseract

PICT

apply(x)

vjp(x, v)

optimized inflow (via gradient)

JAX-CFD

OpenFOAM

INS.jl

deal.II

Warp-NS

···

OpenFOAM

Figure 1 | Overview of Mosaic. Diverse solvers are wrapped behind a uniform containerized interface, enabling cross-solver comparison on shared benchmark tasks involving different physical domains.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

1. Introduction A growing number of physics solvers now advertise gradient support, spanning multiple languages, AD frameworks, and numerical methods. Yet the practical cost of obtaining correct, usable gradients from any of these solvers is opaque. Solvers range from turnkey applications to composable frameworks and low-level kernel toolkits, each demanding different expertise and integration effort. Runtime overhead, gradient accuracy, and numerical conditioning have not been systematically measured, and solver-specific constraints (periodic-only boundary conditions, incompatible penalization schemes, stability limits) surface only after a practitioner has committed to a particular solver. Differentiable solvers enable topology optimization [6], aerodynamic shape optimization [27], optimal control [8], and solver-in-the-loop ML training [30, 51] by backpropagating through a numerical simulation’s forward pass. Differentiating through ODE integrators [15, 46] and renderers [28] is now routine, but PDE solvers are more complex, less standardized, and their gradient properties far less understood. Learned surrogates and physics foundation models offer an alternative path but require perdomain training data, introduce approximation error, and currently cover only a few wellstudied PDE families [11, 42]. Decades of engineering have produced solvers that practitioners trust for forward simulation. What is missing is not the solvers themselves, which exist in abundance, but a standardized way to access, compare, and trust their gradients. We introduce Mosaic, an extensible benchmark suite that tests solver gradients in isolation, on standalone optimization problems without learned components, where finite differences provide ground truth and the solver’s gradient mechanism is the only variable under test. The primary contribution is the benchmark infrastructure itself, not the empirical results from any single instantiation. Researchers can use Mosaic to evaluate which solver fits their problem, build solver-in-the-loop pipelines on top of the uniform interface, and stress-test failure modes before committing to a particular stack. Our contributions are: 1. A benchmark suite and containerized solver interface spanning 14 solvers across four physical domains (heat transfer, structural mechanics, 2D and 3D Navier–Stokes; Table 1). Each solver is wrapped behind a uniform forward/gradient API regardless of language or AD strategy, so researchers can run and compare real solvers without solver-specific expertise. 2. A standardized evaluation protocol measuring gradient accuracy, computational cost, Jacobian conditioning, setup compatibility, and end-to-end optimization convergence. 3. Empirical findings that validate the benchmark by surfacing practically relevant differences: solver choice can determine whether gradient-based workflows succeed or fail, independent of the optimizer, with a negative control that isolates AD bugs from chaotic physics. 4. An open-source release of all solver interfaces, evaluation code, and templates, extensible to new solvers and domains, with continuous re-evaluation designed to keep published results current as solvers improve.

2. Background and related work Adjoint methods and differentiable physics. Adjoint methods for PDE-constrained optimization

date back to Pironneau [43] and Jameson [27], with discrete adjoint formulations following in [19, 40]. Modern AD frameworks now differentiate solver code automatically: JAX source transformation [12], PyTorch operator overloading [41], Julia source-to-source AD via Zygote [26], LLVM-IR-level AD via Enzyme [39], NVIDIA Warp kernel-level AD [37], and

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

tape-based discrete adjoints in dolfin-adjoint and firedrake-adjoint [16]. Hand-derived adjoints are fast but tied to a fixed set of objectives, requiring new derivations whenever the control variable or loss function changes. Algorithmic AD removes this constraint at the price of higher memory or computational cost, a tradeoff whose magnitude has remained largely anecdotal. Mosaic measures it empirically across solver families. Learned surrogates and physics foundation models. Neural operator surrogates [32, 34, 36],

PINNs [47], and physics foundation models [11, 22, 23, 42, 44] are complementary to differentiable solvers but trade accuracy for speed, degrade outside the training distribution, and require solver-generated data in the first place. Mosaic benchmarks the solver layer directly. Existing benchmarks and evaluation efforts. Several benchmarks evaluate PDE solvers or learned

surrogates, but none systematically measures gradient properties across solver backends. APEBench [31] benchmarks neural PDE emulators on forward accuracy without evaluating solver gradients. PDEBench [50] provides datasets of PDE solutions for surrogate training but includes no gradient evaluation. NeuralFluid [33] benchmarks differentiable Navier-Stokes solvers on control tasks and reports gradient runtime, yet does not systematically assess gradient accuracy, conditioning, or setup compatibility across backends. List et al. [35] analyze gradient behavior in unrolled training, focusing on convergence dynamics rather than cross-solver comparison. The PhiFlow benchmark [24] compares AD backends within a single framework, not across independent solver implementations. preCICE [14] standardizes solver coupling but provides no gradient interface. The closest analogue in design philosophy is OpenAI Gym [13], which standardized access to RL environments behind a uniform step/reset interface. Gym ships no RL algorithms, yet the standardized interface enabled a wealth of downstream research by lowering the barrier to entry. Mosaic follows the same model for differentiable solvers: solver backends are simulators behind a standardized gradient interface, and the benchmark ships no ML components itself but provides the infrastructure on which solver-in-the-loop ML research can build.

3. The Mosaic benchmark suite Mosaic organizes evaluation around three concepts. A physical setup is a concrete problem instance: a specific geometry, boundary conditions, and operating regime within a physical domain (e.g., incompressible fluids). A benchmark task is a fixed optimization problem on a specific setup: a control variable, an objective, and a reference solution. A solver backend is the software that evaluates the forward map and (optionally) its gradient, independent of abstraction level. Every solver that supports a task’s physical domain can be compared on it. 3.1 Physical domains and benchmark tasks

Benchmark tasks are representative optimization problems (inflow optimization, topology optimization, inverse problems) that require gradients for efficient solution, are supported by multiple backends for cross-solver comparison, and contain no learned components so solver gradients can be validated in isolation. Tasks are intentionally simpler than state-of-the-art challenges: complexity in the physics would confound the gradient properties we aim to isolate, so any failure can be attributed to the solver’s gradient mechanism rather than to an ill-posed optimization problem. The heat domain (H) additionally serves as a negative control: its smooth, well-conditioned physics means that any gradient failure on H is attributable to an AD bug rather than to chaotic or ill-conditioned dynamics. The four domains form a difficulty progression (Table 1 and Figure 2; full specifications in the supplementary material).

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

Table 1 | Physical domains and optimization tasks. The backend count includes forward-only references; per-solver details are in Table 2.

Domain

Optimization task

Heat transfer (slab)

Backends

Task

Control dim.

Conductivity inversion

128

2 048

12 288

Negative control: smooth physics isolates AD bugs S

Structural mech. (cantilever)

Compliance minimization (SIMP)

Well-studied topology optimization with known optimal F2

Incompr. fluids (2D cylinder)

Inflow optimization for drag min.

Anchor domain; steady flow at Re=20 F3

3D Navier–Stokes (TGV)

Initial condition recovery

Scale test; periodic box isolates time-stepping gradients Observed temperature field

Control: k(x, y)

T = T0 (Dirichlet)

Steady-state solve

Tobs = T(x, y; k * )

Clamped

Control: element densities ρe

F (tip load)

Invert: find k(x, y) that produces Tobs

Objective: min f⊤ u(ρ) s. t.

Objective: min ‖T(k)−Tobs ‖2

∑ veρe ≤ Vf e

(a) H : Steady-state heat conduction — conductivity

(b) S : Cantilever beam — compliance minimization

inversion.

(SIMP). Target: v(T)

Control: v0

Control: inflow profile u(y)

Drag FD

Objective: min FD

Objective: min‖v(T; v0) − vtarget‖2

(d) F3 : 3D periodic box — initial condition recovery.

u(y)

Figure 2 | Benchmark tasks across the four physical domains in Mosaic, ordered by difficulty. Each panel shows the control variable, the physical process, and the optimization objective.

3.2 Design principles

Mosaic consists of three artifacts: a containerized solver interface, an evaluation protocol, and a library of benchmark task templates. Containerization is handled by Tesseract [21], an existing open-source runtime for packaging scientific codes; Mosaic contributes the evaluation methodology, benchmark tasks, and empirical results on top of it. The four domains and 14 solver backends described in this paper are the initial instantiation; the reusable evaluation infrastructure is the lasting contribution, designed to evolve as the community adds solvers, improves configurations, and extends to new domains.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

Solver interface and extensibility. Each solver is a containerized component exposing two entry points, apply(inputs) → outputs and vector_jacobian_product(inputs, cotangents) → vjp.

The container isolates all dependencies (language, libraries, GPU runtime), so any of the 14 backends can be swapped by changing a single configuration string. Forward-only solvers (OpenFOAM, deal.II) implement apply only; Mosaic provides finite-difference gradients automatically at 𝑂 (𝑁 ) cost, serving as the reference baseline. Because solver APIs differ too widely for a single generic wrapper, each benchmark task defines a task-specific interface schema (inputs, outputs, shapes, and units) that multiple solvers implement. Adding a new backend requires only implementing the task schema and adding a container; the evaluation suite re-runs automatically on every proposed change. Benchmark tasks are templates, not fixed endpoints: each template is a validated starting configuration with known-good settings, so users can modify the physical setup (geometry, Reynolds number, boundary conditions) while inheriting the evaluation protocol and finitedifference reference. Details on contributing and containerization overhead are in Section A. Integration effort per backend is quantified in Table 3; column definitions and full experimental specifications are in Section B.1. 3.3 Evaluation protocol

All metrics are computed uniformly across solvers. The protocol is solver-agnostic: it operates on the containerized interface (forward and VJP) and never inspects solver internals. Setup compatibility. For each solver–task pair, we record whether the solver produces a gradient, fails numerically, or is structurally unable to represent the task (Section 4.4). Gradient accuracy. Central finite differences through each solver’s own forward pass serve as the reference. We evaluate directional derivatives along 𝐾 random perturbation vectors (𝐾=6–20) and report cosine similarity and relative ℓ2 error for all problems. For the two Navier–Stokes problems (F2, F3) we additionally compute the full Jacobian and inspect its singular value spectrum. Performance. Primal and VJP wall-clock time, their ratio, and peak (V)RAM consumption, measured on each solver’s intended hardware target across three to four problem sizes (grid resolutions or cell counts). All timings are averaged over 3 runs after 1 warmup iteration (to pre-populate JIT caches and stabilize memory allocation). Forward accuracy. Resolution sweep against a reference solver, adherence to physical laws where applicable (e.g., solenoidal velocity). Where an analytical solution exists (F2, F3), precision is additionally measured against it. Optimization convergence. We run the optimizers with each solver’s gradients on each benchmark task for a fixed iteration budget (500 iterations for F2 and F3; 2 000 for H; 2 500 for S) and report convergence curves and final objective value. Success is defined as reaching a final objective within 1% of the best solution achieved across all solvers within the budget. 3.4 Included solver backends

The solver set (Table 2) is chosen to maximize diversity along three axes: AD strategy (autodiff, tape-based discrete adjoints, hand-derived/kernel adjoints), numerical method (projection-based, lattice Boltzmann, finite-volume, spectral, finite elements with both direct and iterative linear solvers), and solver abstraction level. These range from self-contained applications (OpenFOAM), through frameworks that handle discretization and AD (FEniCS, PhiFlow, JAX-CFD), to kernel toolkits (NVIDIA Warp) where the Mosaic wrapper is the solver. Abstraction level determines available gradient strategies and per-setup effort (Section C).

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

Table 2 | Solver backends included in the current release. Domain IDs match Table 1.

Solver

Language / AD framework Strategy

Discretization Numerics

GPU Domains

JAX-CFD [30]

Py/JAX

Auto

Projection, semi-impl.

✓

PhiFlow [24]

Py/JAX,TF

Auto+An.

Semi-Lagrangian, Euler ✓

F2, F3

INS.jl [1]

Julia/Zygote

Auto

Projection, RK4

F2, F3

XLB [4]

Py/JAX

Auto

LBM

Streaming

✓

F2, F3

PICT [17]

Py/PyTorch

Auto

PISO, BDF1

✓

F2, F3

Warp-NS [37]

Auto

IPCS, SSP-RK3

✓

F2, F3

Exponax [31]

Py/JAX

Auto

Spectral

ETDRK

✓

OpenFOAM [52] C++

—

PISO/PIMPLE, Euler

F2, F3

FEniCS [16]

Sym.+Auto FE

GMRES+AMG

H, S

Firedrake [48]

Sym.+Auto FE

GMRES+AMG

H, S

JAX-FEM [54]

Py/JAX

Auto+An.

Direct (UMFPACK)

TopOpt.jl [25]

Julia

Analytic

Direct (CHOLMOD)

torch-fem [38]

Py/PyTorch

Auto+An.

Direct (spsolve)

deal.II [3, 5]

C++

—

Direct (spsolve)

✓

H, S S

✓

H H, S

Strategy: Auto = framework AD through discrete ops; An. = hand-derived adjoint of the discrete system; Sym. = adjoint derived at variational level via UFL; combinations indicate mixed strategies. Discr.: FD = finite differences; FV = finite volume; FE = finite elements; LBM = lattice Boltzmann method. Numerics: IPCS = incremental pressure correction scheme; PISO = pressure-implicit splitting of operators; PIMPLE = merged PISO-SIMPLE; BDF1 = 1st-order backward differentiation; SSP-RK3 = strong-stability-preserving RK3; ETDRK = exponential time-differencing Runge–Kutta; Streaming = exact LBM streaming step.

Inclusion and exclusion criteria. A solver is included if it provides a general-purpose reverse-

mode VJP at cost proportional to a single forward solve. Forward-only solvers serve as reference baselines. Excluded solvers and reasons are documented in Section C.

4. Experiments The following experiments serve a dual purpose: they validate that Mosaic’s evaluation protocol surfaces practically relevant differences between solvers, and they provide an initial empirical characterization that practitioners can use to guide solver selection. All experiments ran on an Azure Standard_NC24s_v3 instance (4× V100 16 GiB; full spec in Section B.1). Results reflect default solver configurations at the time of the pinned release; all configurations are published and community-improvable via pull request. Reproduction instructions are in Section A. 4.1 Benchmark summary and implementation effort

Table 3 summarizes benchmark results and gradient implementation effort across all solvers; individual metrics are unpacked in the subsections below. Per-domain physical-accuracy validation, including the F2 cylinder forward error and the divergence-RMS, kinetic-energy, and analytic-error sweeps for the NS domains, is reported in Section B.4. Structural incompatibilities exclude some solvers from specific tasks; causes are detailed in Sections C.3 and 4.4. Gradient implementation effort varies substantially among differentiable solvers. We report bespoke lines of code (LOC) as a coarse proxy for integration complexity (an imperfect metric, but one that captures order-of-magnitude differences reproducibly). Warp-NS requires over 800 bespoke lines for the initial-condition VJP alone, a direct consequence of its low-level CUDA

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

Peak (V)RAM

Seconds

MiB

16 GiB

104

Gradient norm 1097 105 103 10

1086 104 102 10

‖∇‖

Wall time

101

102

101

100 10

Rollout steps T

PhiFlow

PICT

Exponax

XLB

Warp-NS

INS.jl

OpenFOAM (FD est.)

102

103 104

Rollout steps T OOM (VRAM) NaN gradient

Figure 3 | VJP cost and stability vs. rollout length on F3.

kernel model where each adjoint pass must be written explicitly. PhiFlow and JAX-CFD each require just 12 shared lines with no per-variable additions, because autodifferentiation through their JAX backends provides gradients at no extra implementation cost. This pattern holds broadly: solvers built on modern AD frameworks (JAX, PyTorch) consistently require zero per-variable extra lines, while hand-written adjoint implementations carry substantially higher overhead. The boundary is not always sharp: INS.jl uses Zygote reverse-mode AD yet still requires over 470 shared lines of gradient-specific code, illustrating that AD framework support does not always translate into turnkey gradients. 4.2 Gradient quality

To assess whether AD gradients are correct, we compare each solver’s gradients against central finite differences at solver-specific optimal perturbation size. All differentiable solvers reach cosine similarity > 0.999 between AD and central finitedifference gradients at their optimal perturbation size 𝜀, confirming AD correctness; per-solver minimum FD errors are reported in the Error (FD) column of Table 3. The optimal 𝜀 shifts by orders of magnitude across solvers and with rollout length and viscosity; see Section B.5 for the full per-domain U-shape curves and discussion. While gradient accuracy is uniformly high, conditioning is not. Figure 5 reports Jacobian conditioning via normalized SVD spectra on the 3D periodic NS problem. Conditioning varies dramatically across solvers at the same resolution, reflecting the numerical method rather than the AD strategy: the same governing equations discretized with lattice Boltzmann versus pressure projection yield radically different spectra. The sharp spectral drop-off in projection-based solvers is attributable to singular vectors associated with non-solenoidal velocity modes, which incompressible flow forbids by construction [10, 20]. This distinction matters for optimization: conditioning predicts which solvers will struggle with second-order methods and long-horizon problems (Sections 4.3 and 4.4). 4.3 Computational cost and scaling

Absolute solver cost varies by orders of magnitude at the same resolution (Table 3; full scaling curves in Section B.3), but speed and generality are entangled: the fastest backends (Exponax, JAX-CFD) support only periodic boundaries and no solid obstacles, while solvers that handle the full F2 drag task (XLB, PICT) are up to ∼200× slower. Cost comparisons are therefore meaningful only within the set of solvers that can represent a given task (Section 4.4).

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

Table 3 | Benchmark results, gradient support, and integration effort per solver. Wall-clock times are means of 3 runs after 1 warmup. Full specifications in Section B.1; initial conditions in Section B.2. Legend below table.

Solver

Wall time Fwd

VJP

Error Ref.

An.

Conv. Grad. support FD

input variables

Navier–Stokes grid (F2), wall time in ms

𝑣0 𝜈

inflow

Integ. effort IO Num. Conf.

Warp-NS [37]

428

1400 1.9e-3 6.2e-4 8.4e-4

◦

0 1631

INS.jl [1]

343

2800 2.8e-3 2.4e-3 1.7e-5

⊥

•

⊥

117 1097

XLB [4]

5200 2.5e-2 2.7e-2 2.1e-5

✓

•

464

PICT [17]

2200

7200 2.4e-3 1.1e-3 1.2e-3

✓

• ⊥

•

329

892

JAX-CFD [30]

108 4.4e-3 4.6e-3 8.3e-4

⊥

•

⊥

286

PhiFlow [24]

713

2200 2.8e-3 2.4e-3 2.2e-5

△

•

262

275

127

437

OpenFOAM [52] 1300

−

2.4e-3

−

Navier–Stokes grid (F3), wall time in ms

𝑣0 𝜈

Warp-NS [37]

486

1500 2.0e-2 2.5e-3 2.5e-5

✓

◦

0 1631

INS.jl [1]

2500 17000 4.1e-2 1.0e-2 9.9e-6

✓

•

117 1097

XLB [4]

43 48000 5.3e-2 1.9e-2 8.0e-6

✓

•

464

PICT [17]

767

3300 1.9e-2 2.5e-3 9.2e-5

✓

• ⊥

329

892

PhiFlow [24]

627

1900 4.0e-2 1.0e-2 1.8e-5

✓

•

262

Exponax [31]

177 2.1e-2 2.5e-3 6.7e-4

✓

•

160

−

275

127

437

OpenFOAM [52] 4400

−

2.5e-3

−

Structural mechanics (S), wall time in s

𝜌

deal.II [5]

166

−

⊥

408

344

147

TopOpt.jl [25]

0.0

−

3.0e-5

✓

◦

124

155

Firedrake [48]

0.0

−

3.0e-5

✓

◦

221

137

FEniCS [16]

140

0.0

−

3.0e-5

✓

◦

222

JAX-FEM [54]

0.0

−

3.0e-5

✓

•

150

133

381

263

Heat transfer (H), wall time in s

𝜌 src

deal.II [5]

9.4

−

⊥ ⊥

FEniCS [16]

0.0

−

2.2e-5

✓

◦

513

torch-fem [38]

0.0

−

2.2e-5

✓

•

268

122

JAX-FEM [54]

4.5

9.8

0.0

−

2.2e-5

✓

•

211

Firedrake [48]

0.0

−

2.2e-5

✓

◦

420

102

Error: Ref. = vs. reference solver; An. = vs. analytical solution (F2/F3); FD = vs. finite differences; − = not applicable. Conv.: ✓ = converged; △ = partial; × = failed; ? = not attempted; ⊥ = structurally excluded (Section C.3). Grad. support: • = native AD; ◦ = implemented by wrapper; × = incorrect; ? = not obtained; ⊥ = structurally blocked. Integ. effort (bespoke LOC): IO = input/output marshalling; Num. = numerics & gradient code; Conf. = solver configuration. PhiFlow partially converges on F2 (△): drag is reduced but the optimizer stalls before reaching the XLB/PICT solution.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

IC recovery error

IC divergence

Optimization loss

100

−2

10−1

10−5

10−2 10

10−8

−3

102

Gradient evaluations

L-BFGS+proj

INS.jl

Adam

L-BFGS

100

101

102

PICT

102

XLB

Figure 4 | F3 IC recovery for selected solvers and different optimizers; full version in Figure 21. Exponax σi / σ1

PhiFlow

steeper ↓ = larger κ (worse conditioning)

10−3 10−6 10−9

10−3

INS.jl σi / σ1

XLB

100

Warp-NS

PICT

100

10−3 10−6 10−9

10−3 10−6

500

1000

1500

Mode index i low ν, T=0.5s

500

1000

1500

Mode index i high ν, T=0.5s

500

1000

1500

Mode index i low ν, T=1s

low ν, T=2s

Figure 5 | Normalized Jacobian singular value spectra (𝜎𝑖 /𝜎0 ) for six solvers on the 3D periodic NS problem at four rollout lengths. Exponax and XLB maintain a relatively flat spectrum across all rollout lengths; INS.jl and PhiFlow develop very steep spectra, indicating large condition numbers 𝜅 = 𝜎1 /𝜎𝑛 .

At larger resolutions and longer rollouts, memory becomes the binding constraint before gradient accuracy does: several solvers exceed the 16 GB GPU memory limit and fail to produce gradients entirely (Figure 3). Among solvers that remain within budget, the cost of obtaining a gradient depends on the AD strategy. Source-transformation AD (JAX-based solvers) maintains a nearly constant VJP/forward ratio across problem sizes, while tape-based adjoints (FEniCS, Firedrake) scale more steeply (Section B.3), making AD strategy a first-order predictor of gradient cost at scale. All timings are measured on V100 GPUs; relative rankings may differ on other hardware, particularly between GPU and CPU solvers. 4.4 Optimization task performance

All solvers share the same hyperparameters and iteration budget per task (settings in Section B.6), so differences in convergence are attributable to the solver’s gradient alone. On the negative-control domains (H, S), all solvers converge to the same optimum, confirming that smooth, well-conditioned physics isolates AD bugs from optimization failure. On F3 (IC recovery), we report the relative ℓ2 distance between the recovered and true initial states (IC error) alongside the optimization loss. Adam stagnates above 40% IC error across all solvers (Figure 4): the loss decreases, but the recovered IC differs noticeably from the target. L-BFGS converges to substantially lower loss and below 6% IC error in fewer iterations. Spectrally projecting the

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

gradient onto the divergence-free subspace before each L-BFGS update (L-BFGS+proj) drops the IC error below 0.5%, because the projection prevents non-solenoidal components from corrupting the curvature model. The 2D drag minimization (F2) is the hardest deployment test, and also the task where structural incompatibilities thin the solver field most: only three of seven backends support the non-periodic, obstacle-laden setup required for drag optimization. Of those, XLB and PICT converge to consistent drag reductions of around 60%; PhiFlow partially converges, reducing drag by ∼35% before stalling. The remaining backends are excluded for structural reasons (periodic BCs, penalization incompatibilities; Section C.3). This attrition itself is a key finding: on the most practically demanding task, solver selection is constrained by setup compatibility before gradient quality enters the picture.

5. Discussion and conclusion Gradient accuracy, computational cost, Jacobian conditioning, and setup compatibility all vary across solvers, across domains, and across operating regimes within a single domain. Yet these differences do not affect optimization outcomes uniformly. Based on the four benchmark tasks evaluated here, we highlight five practitioner-facing observations. Gradient quality is often good enough; the real bottlenecks lie elsewhere. On every task where

they can run, all solvers that produce gradients converge to equivalent optima despite order-ofmagnitude differences in accuracy and conditioning. While more complex tasks may reveal cases where gradient quality becomes limiting, the pattern on these benchmarks suggests that conditioning is likely a better predictor of training stability than pointwise gradient accuracy. Rollout length is the critical scaling parameter. Jacobian conditioning degrades with rollout

length across all solvers (Figure 5), with projection-based solvers deteriorating faster than LBM or spectral methods. Memory scales likewise, and several solvers exceed GPU limits before gradient quality becomes relevant.

On well-conditioned problems, pick the fastest solver. On the negative-control domains (H, S),

all solvers produce equivalent results regardless of AD strategy or numerical method. Because gradient quality is not a differentiator in these regimes, the selection criterion reduces to forward-solve speed and integration effort.

Reverse-mode AD offers the best effort-to-capability tradeoff. Hand-derived adjoints can be faster

than the forward solve (Table 3), but each new control variable or objective requires additional implementation work. Reverse-mode AD through modern frameworks (JAX, PyTorch) produces gradients accurate enough for optimization at a fraction of the implementation cost.

What Mosaic enables beyond this evaluation. Mosaic’s uniform interface makes each backend a

drop-in differentiable simulator for solver-in-the-loop or hybrid physics-ML workflows, and the template mechanism lets the benchmark grow with the community. Limitations. The evaluation rests on three assumptions: that central finite-difference gradients

at optimal perturbation size are a reliable reference (Figure 14), that out-of-the-box solver configurations are representative of practitioner experience, and that standalone optimization tasks are a meaningful proxy for gradient quality in broader pipelines. Gradient quality may differ on untested setups; the template mechanism is the mitigation. We benchmark solver gradients in isolation, without learned components, so hybrid pipelines may introduce failure modes not captured here. Results are pinned to specific solver versions; upstream updates may change gradient behavior. Future work. Natural extensions include additional domains (multiphysics, compressible flow),

solver backends (Enzyme/LLVM-level AD), and hybrid physics-ML tasks. The infrastructure is in place, and we invite community contributions.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

References [1] S. D. Agdestein, S. Ciarella, B. Sanderse, and R. Hoekstra. IncompressibleNavierStokes.jl: Incompressible Navier–Stokes solver in Julia, 2024. URL https://github.com/agdestein/IncompressibleNavierStokes. jl. [2] P. Angot, C.-H. Bruneau, and P. Fabrie. A penalization method to take into account obstacles in incompressible viscous flows. Numerische Mathematik, 81(4):497–520, 1999. doi: 10.1007/s002110050401. [3] D. Arndt, W. Bangerth, M. Bergbauer, M. Feder, M. Fehling, J. Heinz, T. Heister, L. Heltai, M. Kronbichler, M. Maier, P. Munch, J.-P. Pelteret, B. Turcksin, D. Wells, and S. Zampini. The deal.II library, version 9.5. Journal of Numerical Mathematics, 31(3):231–246, 2023. doi: 10.1515/jnma-2023-0089. [4] M. Ataei and H. Salehipour. XLB: A differentiable massively parallel lattice Boltzmann library in Python. Computer Physics Communications, 2024. URL https://arxiv.org/abs/2311.16080. [5] W. Bangerth, R. Hartmann, and G. Kanschat. deal.II—A general-purpose object-oriented finite element library. ACM Transactions on Mathematical Software, 33(4), 2007. doi: 10.1145/1268776.1268779. [6] M. P. Bendsøe and O. Sigmund. Topology Optimization: Theory, Methods, and Applications. Springer, Berlin, 2003. doi: 10.1007/978-3-662-05086-6. [7] D. P. Bertsekas. Projected newton methods for optimization problems with simple constraints. SIAM Journal on Control and Optimization, 20(2):221–246, 1982. doi: 10.1137/0320018. [8] T. R. Bewley, P. Moin, and R. Temam. DNS-based predictive control of turbulence: An optimal benchmark for feedback algorithms. Journal of Fluid Mechanics, 447:179–225, 2001. doi: 10.1017/S0022112001005821. [9] D. A. Bezgin, A. B. Buhendwa, and N. A. Adams. JAX-Fluids: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows. Computer Physics Communications, 282: 108527, 2022. doi: 10.1016/j.cpc.2022.108527. [10] H. Bhatia, G. Norgard, V. Pascucci, and P.-T. Bremer. The Helmholtz–Hodge decomposition—A survey. IEEE Transactions on Visualization and Computer Graphics, 19(8):1386–1404, 2013. [11] C. Bodnar, W. P. Bruinsma, A. Lucic, M. Stanley, A. Vaughan, J. Brandstetter, P. Garvan, M. Riechert, J. A. Weyn, H. Dong, J. K. Gupta, K. Thambiratnam, A. T. Archibald, C.-C. Wu, E. Heider, M. Welling, R. E. Turner, and P. Perdikaris. A foundation model for the Earth system. Nature, 641:1180–1187, 2025. doi: 10.1038/s41586-025-09005-y. [12] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang. JAX: Composable transformations of Python+NumPy programs, 2018. URL https://github.com/jax-ml/jax. [13] G. Brockman, V. Cheung, L. Pettersson, J. Schneider, J. Schulman, J. Tang, and W. Zaremba. OpenAI Gym, 2016. URL https://arxiv.org/abs/1606.01540. [14] H.-J. Bungartz, F. Lindner, B. Gatzhammer, M. Mehl, K. Scheufele, A. Shukaev, and B. Uekermann. preCICE—A fully parallel library for multi-physics surface coupling. Computers & Fluids, 141:250–258, 2016. doi: 10.1016/j.compfluid.2016.04.003. [15] R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31, 2018. [16] P. E. Farrell, D. A. Ham, S. W. Funke, and M. E. Rognes. Automated derivation of the adjoint of high-level transient finite element programs. SIAM Journal on Scientific Computing, 35(4):C369–C393, 2013. doi: 10.1137/120873558. [17] A. Franz, H. Wei, L. Guastoni, and N. Thuerey. PICT – a differentiable, GPU-accelerated multi-block PISO solver for simulation-coupled learning tasks in fluid dynamics. Journal of Computational Physics, 544:114433, 2025. doi: 10.1016/j.jcp.2025.114433. [18] Y. Fu and G. J. Kennedy. Quasi-Newton corrections for compliance and natural frequency topology optimization problems. Structural and Multidisciplinary Optimization, 66:176, 2023. doi: 10.1007/ s00158-023-03630-9.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

[19] M. B. Giles and N. A. Pierce. An introduction to the adjoint approach to design. Flow, Turbulence and Combustion, 65:393–415, 2000. [20] J.-L. Guermond, P. Minev, and J. Shen. An overview of projection methods for incompressible flows. Computer Methods in Applied Mechanics and Engineering, 195(44–47):6011–6045, 2006. [21] D. Häfner and A. Lavin. Tesseract core: Universal, autodiff-native software components for simulation intelligence. Journal of Open Source Software, 10(111):8385, 2025. doi: 10.21105/joss.08385. [22] Z. Hao, C. Su, S. Liu, J. Berner, C. Ying, H. Su, A. Anandkumar, J. Song, and J. Zhu. DPOT: Autoregressive denoising operator transformer for large-scale PDE pre-training. In Proceedings of the 41st International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2024. URL https://proceedings.mlr.press/v235/hao24d.html. [23] M. Herde, B. Raonić, T. Rohner, R. Käppeli, R. Molinaro, E. de Bézenac, and S. Mishra. Poseidon: Efficient foundation models for PDEs. In Advances in Neural Information Processing Systems, 2024. URL https://proceedings.neurips.cc/paper_files/paper/2024/hash/ 84e1b1ec17bb11c57234e96433022a9a-Abstract-Conference.html. [24] P. Holl, V. Koltun, and N. Thuerey. ΦFlow: Differentiable simulations for machine learning. In Proceedings of the 41st International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2024. URL https://proceedings.mlr.press/v235/holl24a.html. [25] Y. Huang and M. Tarek. TopOpt.jl: Truss and continuum topology optimization, interactive visualization, automatic differentiation and more. In Proceedings of the 14th World Congress of Structural and Multidisciplinary Optimization (WCSMO-14), 2021. URL https://github.com/JuliaTopOpt/TopOpt.jl. [26] M. Innes. Don’t unroll adjoint: Differentiating SSA-form programs. arXiv preprint arXiv:1810.07951, 2018. URL https://arxiv.org/abs/1810.07951. [27] A. Jameson. Aerodynamic design via control theory. Journal of Scientific Computing, 3(3):233–260, 1988. doi: 10.1007/BF01061285. [28] H. Kato, D. Beker, M. Morariu, T. Ando, T. Matsuoka, W. Kehl, and A. Gaidon. Differentiable rendering: A survey. arXiv preprint arXiv:2006.12057, 2020. URL https://arxiv.org/abs/2006.12057. [29] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. URL https://arxiv.org/abs/1412.6980. [30] D. Kochkov, J. A. Smith, A. Alieva, Q. Wang, M. P. Brenner, and S. Hoyer. Machine learning–accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, 118(21):e2101784118, 2021. doi: 10.1073/pnas.2101784118. [31] F. Koehler, S. Niedermayr, R. Westermann, and N. Thuerey. APEBench: A benchmark for autoregressive neural emulators of PDEs. In Advances in Neural Information Processing Systems, 2024. URL https: //arxiv.org/abs/2411.00180. [32] N. Kovachki, Z. Li, B. Liu, K. Azizzadenesheli, K. Bhattacharya, A. Stuart, and A. Anandkumar. Neural operator: Learning maps between function spaces with applications to PDEs. Journal of Machine Learning Research, 24(89):1–97, 2023. [33] Y. Li, Y. Sun, P. Ma, E. Sifakis, T. Du, B. Zhu, and W. Matusik. NeuralFluid: Neural fluidic system design and control with differentiable simulation. In Advances in Neural Information Processing Systems, 2024. URL https://arxiv.org/abs/2405.14903. [34] Z. Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, and A. Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021. URL https://arxiv.org/abs/2010.08895. [35] B. List, L.-W. Chen, K. Bali, and N. Thuerey. Differentiability in unrolled training of neural physics simulators on transient dynamics. Computer Methods in Applied Mechanics and Engineering, 2024. URL https://arxiv.org/abs/2402.12971. [36] L. Lu, P. Jin, G. Pang, Z. Zhang, and G. E. Karniadakis. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):218–229, 2021. doi: 10.1038/s42256-021-00302-5.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

[37] M. Macklin. Warp: Differentiable spatial computing for Python. In ACM SIGGRAPH 2024 Courses, pages 1–147. ACM, 2024. doi: 10.1145/3664475.3664543. [38] N. Meyer. torch-fem: Differentiable finite elements for PyTorch, 2024. URL https://github.com/ Meyer-Nils/torch-fem. [39] W. S. Moses and V. Churavy. Instead of rewriting foreign code for machine learning, automatically synthesize fast gradients. In Advances in Neural Information Processing Systems, volume 33, pages 12472–12485, 2020. [40] S. Nadarajah and A. Jameson. A comparison of the continuous and discrete adjoint approach to automatic aerodynamic optimization. In 38th AIAA Aerospace Sciences Meeting and Exhibit, 2000. doi: 10.2514/6.2000-667. [41] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, volume 32, 2019. [42] J. Pathak, S. Subramanian, P. Harrington, S. Raja, A. Chattopadhyay, M. Mardani, T. Kurth, D. Hall, Z. Li, K. Azizzadenesheli, P. Hassanzadeh, K. Kashinath, and A. Anandkumar. FourCastNet: A global data-driven high-resolution weather forecasting model. arXiv preprint arXiv:2202.11214, 2022. URL https://arxiv.org/abs/2202.11214. [43] O. Pironneau. On optimum design in fluid mechanics. Journal of Fluid Mechanics, 64(1):97–110, 1974. doi: 10.1017/S0022112074002023. [44] I. Price, A. Sanchez-Gonzalez, F. Alet, T. Ewalds, A. El-Kadi, J. Stott, S. Mohamed, P. Battaglia, R. Lam, and M. Willson. Probabilistic weather forecasting with machine learning. Nature, 637:1038–1044, 2024. doi: 10.1038/s41586-024-08252-9. [45] Y. Qian, D. d’Humières, and P. Lallemand. Lattice BGK models for Navier–Stokes equation. Europhysics Letters, 17(6):479–484, 1992. [46] C. Rackauckas, Y. Ma, J. Martensen, C. Warner, K. Zubov, R. Supekar, D. Skinner, A. Ramadhan, and A. Edelman. Universal differential equations for scientific machine learning. arXiv preprint arXiv:2001.04385, 2021. [47] M. Raissi, P. Perdikaris, and G. E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686–707, 2019. doi: 10.1016/j.jcp.2018.10.045. [48] F. Rathgeber, D. A. Ham, L. Mitchell, M. Lange, F. Luporini, A. T. T. McRae, G.-T. Bercea, G. R. Markall, and P. H. J. Kelly. Firedrake: Automating the finite element method by composing abstractions. ACM Transactions on Mathematical Software, 43(3):1–27, 2016. doi: 10.1145/2998441. [49] K. Svanberg. The method of moving asymptotes — a new method for structural optimization. International Journal for Numerical Methods in Engineering, 24(2):359–373, 1987. doi: 10.1002/nme.1620240207. [50] M. Takamoto, T. Praditia, R. Leiteritz, D. MacKinlay, F. Alesiani, D. Pflüger, and M. Niepert. PDEBench: An extensive benchmark for scientific machine learning. Advances in Neural Information Processing Systems, 35:1596–1611, 2022. [51] K. Um, R. Brand, Y. R. Fei, P. Holl, and N. Thuerey. Solver-in-the-loop: Learning from differentiable physics to interact with iterative PDE-solvers. Advances in Neural Information Processing Systems, 33: 6111–6122, 2020. [52] H. G. Weller, G. Tabor, H. Jasak, and C. Fureby. A tensorial approach to computational continuum mechanics using object-oriented techniques. Computers in Physics, 12(6):620–631, 1998. doi: 10.1063/1. 168744. [53] G. D. Weymouth and B. Font. WaterLily.jl: A differentiable and backend-agnostic Julia solver for incompressible viscous flow around dynamic bodies. Computer Physics Communications, 315:109748, 2025. doi: 10.1016/j.cpc.2025.109748. [54] T. Xue, S. Liao, Z. Gan, C. Park, X. Xie, W. K. Liu, and J. Cao. JAX-FEM: A differentiable GPU-accelerated

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

3d finite element solver for automatic inverse design and mechanistic data science. Computer Physics Communications, 291:108802, 2023. doi: 10.1016/j.cpc.2023.108802.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

A. Software and reproducibility A.1 License, access, and reproduction

Mosaic is released under the Apache 2.0 license (individual solver backends retain their upstream licenses). The repository, documentation, and reproduction instructions are available at https: //github.com/pasteurlabs/mosaic. The paper corresponds to the tag v0.1+paper-repro which contains all code, plotting scripts, solver wrappers, and a pinned dependency lockfile sufficient for full reproduction. Two reproduction paths are supported. Figures can be regenerated from the result archive linked in the repository README without re-running any solver. Alternatively, the full experiment suite can be re-run from scratch; this requires Docker and, for GPU solvers, an NVIDIA GPU with the NVIDIA Container Toolkit. The lockfile pins every transitive dependency to the exact versions used for the paper, and solver containers are built on-demand by the harness. Step-by-step instructions for both paths are in the repository README. A.2 Architecture

Each solver backend is a Tesseract [21]: a Docker container that exposes a typed forward map (apply), shape inference (abstract_eval), and optionally a VJP (vector_jacobian_product) over Pydantic schemas. The container boundary isolates solver-specific dependencies (JAX, PyTorch, Julia, C++, Warp) from one another and from the evaluation harness. tesseract-jax bridges the container interface into JAX’s jit/grad transform chain, so the harness can differentiate through any backend regardless of its native AD framework using a single code path. Each backend is defined by three files: a Python module implementing the entry points above, a YAML configuration declaring solver metadata (name, AD strategy, backend framework, plot style), and a requirements file listing dependencies. The evaluation harness discovers solvers automatically from their configuration metadata and runs five benchmark suites (forward accuracy, gradient quality, computational cost, optimization convergence, and initial-condition visualization) without solver-specific code. A.3 Extensibility

Contributing a new solver requires implementing the three files described above for an existing domain. The harness discovers the backend automatically from its configuration metadata; no registration step or evaluation-code changes are needed. A step-by-step tutorial (adding an LBM solver to the fluid domain) is included in the repository documentation. New benchmark domains can be scaffolded from built-in templates and require interface schemas, at least one solver backend, and a problem configuration defining initial conditions, error metrics, and suite defaults. Existing solver configurations can also be improved via pull requests; CI re-runs the full evaluation suite and posts a before/after comparison. A.4 Maintenance and versioning

We maintain the evaluation protocol, the Tesseract interface contract, CI infrastructure, and the paper-pinned solver set (tagged release, kept buildable). New solver contributions, domainspecific task formulations, and upstream solver maintenance are community-owned. If an upstream solver breaks due to a dependency update, we document version pins and accept community fixes but do not patch upstream code ourselves. Each release is tagged with a semantic version and accompanied by a pinned dependency lockfile. Continuous integration runs the benchmark suite on every push to main that touches a solver or the harness, committing results to a dedicated branch with a machine-readable diff against the previous baseline. Tagged releases run the full suite and publish a result archive

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

as a GitHub release artifact. This keeps published results current as solvers evolve and new backends land.

B. Extended experimental results B.1 Benchmark specifications Integration effort. The IO, Num., and Conf. columns in Table 3 report bespoke lines of code

by category. IO covers serialization and file-based data exchange with external solver processes. Num. is problem-specific numerical code (assembly, time-stepping, boundary conditions, and gradient plumbing). Conf. covers solver initialization and mesh setup. Glue code (Tesseract schema definitions and entry-point boilerplate) is excluded; it consistently amounts to 100– 400 LOC across all backends. Error column definitions. Ref.: ∥fsolver − fref ∥ 2 /∥fref ∥ 2 , where fref is the forward state from the

reference solver (OpenFOAM for F2/F3; deal.II for S/H). This is a forward accuracy metric, not a gradient comparison. An.: ∥fsolver −fTGV ∥ 2 /∥fTGV ∥ 2 , where fTGV is the TGV closed-form forward solution (F2/F3 only). FD: minimum over 𝜀 ∈ {10−6, . . . , 10−1 } of ∥gsolver − gFD,𝜀 ∥ 2 /∥gFD,𝜀 ∥ 2 , where gFD,𝜀 is a central finite-difference gradient estimate at step size 𝜀. Hardware and timing protocol. All experiments ran on an Azure Standard_NC24s_v3 instance (24 vCPUs, 448 GiB RAM, 4× NVIDIA Tesla V100 16 GiB, Intel Xeon E5-2690 v4 @ 2.60 GHz). Wall times use one warmup call (excluded, absorbs JIT compilation and CUDA kernel caching) followed by three timed trials; the reported value is the mean. The timed window covers hostside serialization, the Tesseract HTTP round-trip, solver compute, and response deserialization; container startup is excluded. F2 (2D NS, grid). Periodic domain [0, 2𝜋) 2 , uniform Cartesian grid, Taylor–Green vortex IC.

Fwd/VJP: 64 × 64 grid (4 096 cells), 𝜈 = 0.01, Δ𝑡 = 0.01, 100 steps. Error (Ref., An.): agreement experiment, 64 × 64 grid, 𝜈 = 0.01, Δ𝑡 = 0.05, 20 steps; Ref. vs. fine JAX-CFD reference (Δ𝑡 = 0.01, 100 steps), An. vs. TGV closed-form u(𝑡) = u(0)𝑒 −2𝜈𝑡 . Error (FD): finite-difference gradient check, 16 × 16 grid, multimode IC, 𝜈 = 0.001, Δ𝑡 = 0.05, 20 steps, 20 random directions. Convergence: drag-reduction optimization, 32 × 32 grid, channel domain [0, 1] 2 , cylinder at (0.5, 0.5) radius 0.05, Re = 20 (𝜈 = 0.0025), Δ𝑡 = 0.02, 400 steps, Adam (𝜂 = 5 × 10−4 ), up to 500 iterations. An L-BFGS variant with the same physics is reported in Section B.6. F3 (3D NS, grid). Periodic domain [0, 2𝜋) 3 , uniform Cartesian grid, TGV IC. Fwd/VJP: 163

grid (4 096 cells), 𝜈 = 0.01, Δ𝑡 = 0.01, 50 steps. Error (Ref., An.): forward agreement, 163 grid, 𝜈 = 0.01, Δ𝑡 = 0.01, 50 steps; Ref. vs. fine Exponax reference (Δ𝑡 = 0.002, 250 steps), An. vs. TGV closed-form. Error (FD): finite-difference gradient check, 𝜈 = 0.001, Δ𝑡 = 0.05, 10 steps, 10 random directions. Convergence: IC recovery, 𝜈 = 0.01, Δ𝑡 = 0.02, 100 steps, Adam (𝜂 = 10−3 ), up to 500 iterations. S (Structural mechanics, mesh). Cantilever beam [0, 2] × [0, 1] × [0, 1], HEX8 elements, SIMP

penalization (𝑝 = 3, 𝐸 max = 70 000 MPa, 𝜈 mat = 0.3, 𝑥 min = 10−3 ). Fwd/VJP: 128 × 2 × 64 mesh (≈16k elements), uniform density 𝜌 0 = 0.5, distributed unit load 𝐹 = 1. Error (FD): 8 × 2 × 4 mesh (64 elements), corner point load, 6 random directions. Convergence: compliance minimization, 16 × 2 × 8 mesh (256 elements), volume fraction 0.5, Adam (𝜂 = 5 × 10−2 ), up to 2 500 iterations. H (Heat transfer, mesh). Quasi-2D heated slab [0, 2] × [0, 1], HEX8 elements (single layer), SIMP

penalization (𝑝 = 3, 𝑘 max = 1, 𝑘 min /𝑘 max = 10−3 ). Fwd/VJP: 256 × 128 × 1 mesh (≈33k elements), uniform conductivity 𝑘 0 = 0.5, unit distributed heat source. Error (FD): 8 × 4 × 1 mesh (32 elements), 6 random directions. Convergence: conductivity recovery from two-Gaussian target, 16 × 8 × 1 mesh, Adam (𝜂 = 10−2 ), up to 2 000 iterations.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

B.2 Initial conditions

Each benchmark domain uses a fixed set of initial conditions, shared across all solvers. Incompressible NS (Figure 6). Taylor–Green vortex (TGV, 2D): 𝑢 = sin(𝑥) cos(𝑦), 𝑣 =

− cos(𝑥) sin(𝑦) on [0, 2𝜋) 2 . Admits the analytic solution u(𝑡) = u(0)𝑒 −2𝜈𝑡 , enabling exact error measurement without a reference solver. Multimode: solenoidal field constructed from a spectral-space stream function with energy concentrated in a ring at wavenumber 𝑘 = 2 (𝜎𝑘 = 0.5) and random phases; peak speed normalized to 0.3. Uniform inflow and flat inlet profile are spatially constant fields; not shown. Taylor–Green vortex (TGV, 3D): 𝑢 = sin(𝑥) cos(𝑦) cos(𝑧), 𝑣 = − cos(𝑥) sin(𝑦) cos(𝑧), 𝑤 = 0 on [0, 2𝜋) 3 . Evolves into a vortexdominated turbulent state under the NS dynamics. Arnold–Beltrami–Childress (ABC): 𝑢 = 𝐴 sin(𝑧) + 𝐶 cos(𝑦), 𝑣 = 𝐵 sin(𝑥) + 𝐴 cos(𝑧), 𝑤 = 𝐶 sin(𝑦) + 𝐵 cos(𝑥) with 𝐴 = 𝐵 = 𝐶 = 1 on [0, 2𝜋) 3 . A steady Euler solution with chaotic particle trajectories. TGV (2D)

Multimode (2D)

TGV (3D)

ABC (3D)

Figure 6 | Initial conditions for the F2 and F3 domains (vorticity; 3D shown as 𝑧=0 slice).

Structural mechanics (Figure 7). All ICs are material density fields over the beam mesh. Uni-

form: 𝜌 0 = 0.5 everywhere. Random: Gaussian noise centered at 𝜌 0 = 0.5 (𝜎 = 0.3), clipped to [0.05, 0.95]. Two density bumps: two Gaussian pillars (𝜌 peak = 0.95, 𝜎 = 0.12 𝐿) at 𝑥 = 0.35 𝐿𝑥 and 𝑥 = 0.75 𝐿𝑥 , on a soft background (𝜌 bg = 0.1). Uniform density

Random density

Two density bumps

Figure 7 | Initial conditions for the structural mechanics domain.

Heat transfer (Figure 8). All ICs are conductivity or source fields over the thermal mesh. Uni-

form conductivity: 𝑘 0 = 0.5 everywhere. Random conductivity: Gaussian noise centered at 𝑘 0 = 0.5 (𝜎 = 0.3), clipped to [0.05, 0.95]. Gaussian source: single heat source centered at (0.5 𝐿𝑥 , 0.5 𝐿𝑦 ) with width 𝜎 = 0.2 min(𝐿𝑥 , 𝐿𝑦 ). Two-Gaussian conductivity: groundtruth field for the conductivity recovery task; two Gaussian peaks at (0.3 𝐿𝑥 , 0.5 𝐿𝑦 ) and (0.7 𝐿𝑥 , 0.5 𝐿𝑦 ) with width 𝜎 = 0.15 min(𝐿𝑥 , 𝐿𝑦 ). B.3 Cost

Figures 9 and 10 show forward time, VJP time, and VJP/forward ratio vs. element count for all four domains. The x-axis is total element count, parameterized by a single integer 𝑁 : F2 uses an 𝑁 × 𝑁 uniform Cartesian grid (𝑁 ∈ {64, 128, 192, 256}, giving 4k–66k cells); F3 uses an 𝑁 × 𝑁 × 𝑁 Cartesian grid (𝑁 ∈ {16, 32, 48, 64}, giving 4k–262k cells); S uses an 𝑁 × 2 × ⌊𝑁 /2⌋ hexahedral mesh (𝑁 ∈ {8, . . . , 3200}, giving 64–10M elements); H uses an 𝑁 × ⌊𝑁 /2⌋ × 1 hex

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

Uniform conductivity

Random conductivity

Gaussian source

Figure 8 | Initial conditions for the heat transfer domain.

mesh (𝑁 ∈ {16, . . . , 4500}, giving 128–10M elements). Table 3 wall times are at 𝑁 = 64 (F2), 𝑁 = 16 (F3), 𝑁 = 128 (S), 𝑁 = 256 (H). The fixed RPC overhead (host serialization and HTTP round-trip) adds approximately 2–20 ms per call and is visible as a floor for the fastest solvers at small problem sizes. Lines that terminate before the full size range indicate that the solver either raised an exception (typically an out-of-memory error) or exceeded the 1 000 s per-trial wall limit at that problem size (the underlying Tesseract HTTP request has a separate 1 200 s watchdog); all larger sizes for that solver are omitted. Forward times span several orders of magnitude across solver families at identical resolution, reflecting differences in algorithmic complexity and hardware target (CPU vs. GPU). VJP overhead tracks forward cost for source-transformation AD (JAX). For the NS domains the VJP/forward ratio is roughly constant across problem sizes: both forward and backward passes integrate the same 𝑇 timesteps, so their costs scale identically with 𝑁 . For structural and thermal domains the ratio sits close to 1 at moderate sizes because the adjoint of a symmetric elliptic system shares the same stiffness matrix as the forward solve and can reuse its factorization; at the largest problem sizes the adjoint-based solvers (FEniCS, Firedrake) show steeper VJP scaling as the adjoint solve cost begins to dominate.

2D NS Time (s)

Forward time

VJP time

101

102

100

101

10−1

100

100 10−1 10−2

VJP / forward

3D NS Time (s)

DOFs 102

101

103

101

100 10−1

10−2

10−1

JAX-CFD (G) PhiFlow (G)

1k 62k 2 INS.jl (C) XLB (G)

102

101 100

DOFs

PICT (G) Warp-NS (G)

1k 62k 2

Exponax (G)

1k 62k 2

OpenFOAM (C)

Figure 9 | Gradient overhead scaling for the fluid domains (2D NS, top; 3D NS, bottom). Left: forward time; center: VJP time; right: VJP/forward ratio (dashed line marks 1:1). All axes are log-log. (G) / (C) denotes GPU / CPU execution.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

Structural Time (s)

Forward time

VJP time 102

100 10

VJP / forward

103

102 10

100

10−1

−1

DOFs Thermal Time (s)

102

101 10

101

100

10−1

JAX-FEM (G) TopOpt.jl (C)

4k deal.II (C)

DOFs

FEniCS (C)

Firedrake (C)

TorchFEM (G)

Figure 10 | Gradient overhead scaling for the mesh domains (Structural, top; Thermal, bottom). Same layout as Figure 9. deal.II VJP times are omitted (no native adjoint). Source-transformation AD (JAX-FEM) maintains a nearly constant ratio; adjoint-based solvers (FEniCS, Firedrake) show steeper VJP scaling at large problem sizes.

B.4 Physical accuracy

As a basic sanity check, we run agreement and baseline experiments on each NS domain: agreement measures relative ℓ2 error vs. the analytic TGV solution across a 𝜈 sweep at fixed resolution, and baseline measures convergence with resolution at one timestep. All solvers that support the periodic TGV domain pass both checks with errors below 10−2 . F2 cylinder-flow forward accuracy. The F2 domain uses a cylinder obstacle in a channel. Unlike

the periodic TGV, no analytic solution exists, so solvers are compared against the cross-solver consensus (mean over valid solvers). This experiment uses a unit-square domain [0, 1] 2 with a cylinder at (0.5, 0.5), radius 0.1, 𝑁 = 64, 500 steps at Δ𝑡 = 0.01; the physical accuracy sweeps below use the full channel domain [0, 8] × [0, 2]. The F2 drag-optimization convergence experiment (Section B.1) uses the same unit-square domain but a smaller cylinder (radius 0.05, 𝐷 = 0.1), which sets Re = 𝑈 𝐷/𝜈 = 20 at the fixed 𝑈 = 0.5, 𝜈 = 0.0025 used there; both setups are F2 cylinder-channel flows but exercise different aspects of the benchmark. Three solvers (JAX-CFD, INS.jl, Warp-NS) produce invalid results on this geometry due to boundary condition incompatibilities. Figure 11 shows consensus error vs. 𝜈 and final vorticity fields for the four solvers that run successfully. Physical accuracy sweeps. For the two fluid domains, Figures 12 and 13 show resolution 𝑁 ,

viscosity 𝜈, and rollout length against three physical metrics. F2 uses the [0, 8] × [0, 2] cylinderchannel at Re = 20, sweeping 𝑁 ∈ {16, 32, 64, 128}, 𝜈 ∈ {0.001, 0.005, 0.01, 0.05, 0.1}, and up to 100 steps. F3 uses the [0, 2𝜋] 3 TGV box at 𝑁 = 16, sweeping 𝜈 ∈ {0.001, 0.01, 0.05, 0.1} and up to 50 steps. Analytic TGV error is the relative ℓ2 error against the Taylor–Green vortex closed-form solution. Lower is better, with values below 10−2 indicating good accuracy. Divergence RMS measures how strongly the solver violates the incompressibility constraint ∇ ·u = 0. For a well-posed incompressible solver this should be near machine precision (≲ 10−6 ), while values above 10−2 indicate a solver that does not enforce incompressibility exactly (e.g. LBM methods). Kinetic energy is the domain-averaged 12 ∥u∥ 2 . For 2D NS the analytical

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

Final vorticity (ν = 0.01)

Consensus error

0.40

0.35

PhiFlow XLB PICT OpenFOAM

PhiFlow

OpenFOAM

PICT

XLB

0.30

0.25

0.20

05 0.0

10 0.0

20 0.0

0.0

Figure 11 | F2 cylinder-flow forward accuracy: relative ℓ2 error vs. the cross-solver consensus as a function of 𝜈 at 𝑁 = 64, 500 steps (left), and final vorticity fields at 𝜈 = 0.01 (right). Solvers that fail on this geometry are omitted.

ground truth is KE(𝑡) = 14 𝑒 −4𝜈𝑡 , shown as a dashed reference line, and solvers should track it closely. For 3D NS no closed-form solution exists, so all solvers should at least agree with each other; spread at a given parameter value indicates a physical accuracy problem. Solvers that enforce the divergence-free constraint exactly at each step by construction (PhiFlow, Warp-NS, PICT, JAX-CFD, INS.jl, Exponax) sit at or near machine precision regardless of resolution or viscosity. XLB exhibits a compressibility floor in divergence RMS and analytic error that is largely independent of 𝜈 or resolution, consistent with the 𝑂 (Ma2 ) compressibility error inherent to the LBM formulation. For the mesh domains, all structural and thermal solvers agree to machine precision across the full load range, consistent with the expected 𝐶 ∝ 𝐹 2 scaling of linear compliance. B.5 Gradient quality Finite-difference verification. Figure 14 extends the FD verification from Section 4.2 to all four

domains. The U-shaped error curve (truncation error at large 𝜀, floating-point roundoff at small 𝜀) is visible in every domain and for every solver that provides gradients. The optimal 𝜀 shifts noticeably between solver families and between domains: spectral solvers (Exponax) tolerate larger perturbations than projection-based solvers, and structural/thermal solvers operate at a different optimal scale than fluid solvers. All differentiable solvers reach cosine similarity > 0.999 at their optimal 𝜀, confirming gradient correctness across domains. Jacobian singular value spectra. Figure 16 shows normalized singular value spectra 𝜎𝑖 /𝜎0 of the

solver Jacobian for 2D NS across varying viscosity and rollout length; the 3D NS spectra appear in Figure 5 in the main text. Faster spectral decay indicates a more ill-conditioned Jacobian and correspondingly harder gradient-based optimization. Across both domains, increasing rollout length (𝑇 ) consistently steepens the decay, while higher viscosity (𝜈) slightly moderates it. Solver families cluster: spectral solvers retain flatter spectra at longer horizons, while projection-based solvers exhibit sharper decay. The sharp spectral drop-off in projection-based solvers is attributable to singular vectors associated with non-solenoidal velocity modes, which incompressible flow forbids by construction [10, 20]; these higher modes are mostly attributed to divergence and are correspondingly less relevant for gradient-based optimization.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

2D NS — physical accuracy

Analytic error

Divergence RMS 10−2

10−3

10−4

0.23

10−4

10−6

vs N

10−2

vs ν

10−2 10−3 10−4 1

vs steps

Kinetic energy 0.24

0 0 05 10 0. 0.

0.22 1k

10−2

0.25

10−4

0.24

10−6

0.23 1

10−2

10−3

10−4

0 0 05 10 0. 0.

0.24 0.22

10−6

10−4 510

510

Steps JAX-CFD

510

Steps INS.jl

XLB

PICT

Steps Warp-NS

OpenFOAM

PhiFlow

Figure 12 | 2D NS physical accuracy. Each panel sweeps one parameter (resolution, 𝜈 , or steps) on the x-axis against one physical metric on the y-axis.

Gradient quality over rollout length. Figure 17 shows how FD gradient quality evolves with

rollout length on the 2D NS domain. Even solvers with correct gradients at short horizons can degrade at longer rollouts, consistent with the Jacobian conditioning analysis in Figure 5: solvers with rapidly decaying singular value spectra lose useful gradient signal faster. The optimal perturbation size 𝜀 also drifts with horizon, reinforcing that a fixed FD step used for validation can give misleading results at longer rollout lengths. FD U-curves. Figures 18 and 19 show full FD U-curves (relative error vs. perturbation size 𝜀)

at each rollout length, for the 2D and 3D NS domains respectively. Each U-curve has a characteristic minimum: too-small 𝜀 is dominated by floating-point cancellation; too-large 𝜀 by nonlinear truncation error. The minimum shifts rightward and rises as the rollout grows, reflecting increasing Jacobian ill-conditioning. B.6 Optimization Optimizer configurations. All optimization runs use either Adam (optax.adam, default 𝛽 1 =0.9, 𝛽 2 =0.999, 𝜖=10−8 ) or L-BFGS (optax.lbfgs with the default zoom line search and a 10-step

inverse-Hessian memory), with patience-based early stopping on the loss: the run terminates as soon as the loss has not strictly improved over the running minimum for patience consecu-

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

3D NS — physical accuracy

vs N

Analytic error

Divergence RMS 10−2

10−1

0.10

10−3

0.08

10−4 k

vs ν 1

vs steps

0 0 05 10 0. 0.

1 00

0 0 05 10 0. 0.

XLB

0 0 05 10 0. 0.

0.10 5 10

Steps PhiFlow

0.11

10−4 20

0.12

10−3

5 10

0.11

10−2

10−1

0.12

10−3

10−2

Kinetic energy 0.12

Steps PICT

Warp-NS

5 10

Steps Exponax

OpenFOAM

INS.jl

Figure 13 | 3D NS physical accuracy on the triply-periodic TGV domain [0, 2𝜋) 3 . Rows sweep resolution 𝑁 (top), viscosity 𝜈 (middle), and rollout length (bottom). Columns report analytic TGV error, divergence RMS, and kinetic energy. No closed-form kinetic-energy reference exists for 3D NS, so solvers are compared against each other; spread at a given parameter value indicates a physical accuracy discrepancy.

tive iterations, or when the maximum iteration budget is reached. The 3D NS recovery and thermal conductivity convergence plots use gradient evaluations on the x-axis rather than outer iterations: Adam costs one value-and-gradient call per step, while L-BFGS’s zoom line search adds an empirically observed average of ∼3 probes per outer step, so a per-iteration x-axis would understate L-BFGS’s true forward-and-backward cost. Per-task settings: • H (conductivity recovery): Adam 𝜂=10−2 , 2 000 iters, patience 200; L-BFGS, 200 iters, patience 30. • S (topology optimization): Adam 𝜂=5×10−2 , 2 500 iters, patience 100, with a soft volume-fraction penalty of weight 50; L-BFGS, 100 iters, patience 20 (fails to reach Adam’s optimum, see below); MMA via NLopt [49], 200 iters, patience 30, with a hard volume inequality and native box constraints. • F2 (drag minimization, Re=20): Adam 𝜂=5×10−4 , 500 iters, patience 100, with a flow-rate penalty of weight 50; L-BFGS, 50 iters, patience 15 (diverges, see below). • F3 (3D NS IC recovery): Adam 𝜂=10−3 , 500 iters, patience 50; L-BFGS, 100 iters, patience 20. L-BFGS+proj and Adam+proj variants apply a solenoidal projection to the gradient

1 - cos. sim.

Relative FD error

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

2D NS 10

10−2

10−3 10

−4

3D NS 100

−1

−4

10−1

10−6

10−8

Structural

−7

Thermal −1

−1

−2

10−2

10−3

10−4

10−2

−4

10−4

10−6

10−8

10−4

10−2

JAX-CFD

XLB

Exponax

deal.II

Firedrake

PhiFlow

PICT

JAX-FEM

FEniCS

TorchFEM

INS.jl

Warp-NS

TopOpt.jl

Figure 14 | FD gradient verification across all four domains. Top row: relative ℓ2 error vs. perturbation size 𝜀 . Bottom row: subspace cosine similarity vs. 𝜀 . All differentiable solvers achieve cosine similarity > 0.999 at their optimal 𝜀 . Cosine similarity (2D NS) 1 − cos(∇AD, ∇FD)

Rel. FD error

Relative error (2D NS) 10−1 10−2 10−3 10−4 10−4 10−3 10−2 10−1

10−4 10−6 10−8 10−4 10−3 10−2 10−1

100

Perturbation size ε JAX-CFD

PhiFlow

100

Perturbation size ε INS.jl

XLB

PICT

Warp-NS

Figure 15 | FD gradient verification on the 2D NS domain (focused view of the leftmost column of Figure 14). Relative

ℓ2 error (left) and cosine similarity (right) vs. perturbation size 𝜀 , normalized by the RMS magnitude of the initial condition. The U-shaped curve reflects the truncation/roundoff tradeoff; optimal 𝜀 varies by orders of magnitude across solvers.

at each step (see next paragraph) with otherwise identical settings. All variants start from 𝑢 ≡ 0 and minimize the ℓ2 distance to the target final state, averaged over three IC seeds. The seeds vary the ground-truth target field (rand_div_free regenerated with seeds 0, 1, 2), not the optimizer initialization, which is identically zero across seeds; this measures recovery of arbitrary div-free targets from a cold start rather than re-running the same problem three times. Solenoidal gradient projection (F3 variants). The L-BFGS+proj and Adam+proj variants on F3

apply a spectral Helmholtz projection to the gradient before each optimizer update, 𝑔ˆ𝑖df = ˆ 𝑔ˆ𝑖 −𝑘𝑖 (𝑘·𝑔)/|𝑘 | 2 , where 𝑔ˆ is the FFT of the gradient on the periodic box. This restricts the search direction to the divergence-free subspace and is exact to machine precision on the periodic NS state space. The projection is implemented once in the harness and applied uniformly to all six

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

JAX-CFD

PhiFlow

σi / σ1

100 steeper ↓ = larger κ (worse conditioning)

10−3 10−6 10−9

100

10−3 10−6 10−9

XLB σi / σ1

INS.jl

100

PICT

Warp-NS

100

10−3 10−6 10−9

10−3 10−6

10−3 10−6 10−9

120

Mode index i

120

Mode index i

low ν, T=0.5s

high ν, T=0.5s

120

Mode index i low ν, T=1s

low ν, T=2s

Figure 16 | Normalized Jacobian singular value spectra for 2D NS. Each panel shows a different (𝜈,𝑇 ) configuration. Faster decay corresponds to a more ill-conditioned Jacobian and harder gradient optimization. Cosine similarity (best ε) 10−6

2 × 100

101

10−7

10−4 10−5

10−8 10−9 10−10 10−11

101

102

Rollout steps T JAX-CFD

10−3

1 − cosine

3 × 100

‖∇‖

FD relative error (best ε)

Gradient norm Relative FD error

4 × 100

101

102

Rollout steps T PhiFlow

INS.jl

XLB

102

Rollout steps T PICT

Warp-NS

Figure 17 | Gradient quality vs. rollout length (2D NS). Left: gradient norm over rollout length. Middle: best-achievable FD relative error at each horizon. Right: optimal perturbation size 𝜀 . Solvers with poor Jacobian conditioning degrade fastest with increasing rollout length.

fluid solvers; we do not project elsewhere because the divergence-free constraint is specific to the incompressible NS state. 2D drag minimization. Figure 20 shows convergence for the 2D cylinder-flow drag minimization

task. This task is notably difficult to deploy: the channel geometry with advective outflow requires non-periodic pressure boundary conditions that exclude spectral solvers, while Brinkman penalization of the cylinder obstacle is incompatible with LU-factored Poisson projections, collectively ruling out four of the seven F2 backends (detailed in Section C.3). XLB and PICT converge to consistent drag reduction of around 60% at Re = 20; PhiFlow partially converges, reducing drag by ∼35% before stalling. We also attempted L-BFGS on this task. It failed universally: XLB diverged to NaN within 50 iterations (drag coefficients blowing up from 0.083 to over 8) and PICT produced no converged result. The failure is structural: the inflow-to-drag map is strongly non-convex (advection-driven, with separation and re-attachment regimes), so secant pairs 𝑠𝑘 = 𝑥𝑘+1 − 𝑥𝑘 taken across iterations sample regions where the local Hessian differs in sign, violating the positive-curvature condition and causing the limited-memory model to blow up. This contrasts with the quasi-static problems in Figures 23 and 24, where L-BFGS

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026 T=5

T = 10

Rel. FD error

10−2 10−3 10−4

−3

10−3

T = 40 10

−2

10−3 10−4

10−5 10−3 10−2 10−1 100

10−3 10−2 10−1 100

T = 80

T = 160

T = 320

10−1

10−2

10−3

10−4 10−3 10−2 10−1 100

ε JAX-CFD

−1

10−2

10−4

10−5 10−3 10−2 10−1 100

Rel. FD error

T = 20

−2

10−4

10−5

10−3

−4

10−4 10−3 10−2 10−1 100

PhiFlow

10−3 10−2 10−1 100

INS.jl

XLB

PICT

Warp-NS

Figure 18 | FD U-curves for F2 (2D NS): relative FD error vs. perturbation size 𝜀 at each rollout length 𝑇 . T = 10

T = 20

Rel. FD error

100 10

−2

10−2

10−3

10−4

10−5

10−4

10−310−210−1 100

ε PhiFlow

T = 80

T = 160 100

10−1

10−1 −2

T = 40 10−1

10−1

10−2

10−3

10−4

10−310−210−1 100

ε Warp-NS

INS.jl

XLB

PICT

10−310−210−1 100

ε Exponax

Figure 19 | FD U-curves for F3 (3D NS): relative FD error vs. perturbation size 𝜀 at each rollout length 𝑇 .

converges reliably; the determining factor is whether the landscape is stationary enough to sustain a coherent curvature model across iterations. 3D NS IC recovery. Figure 21 compares three optimizer configurations on the 3D NS IC recovery

task ([0, 2𝜋) 3 TGV box, 163 grid). Each run starts from a constant zero-velocity field and minimizes the ℓ2 distance to a target final state. We report the IC error as the relative ℓ2 distance between the recovered and true initial velocity fields, ∥u∗0 − u0 ∥ 2 /∥u0 ∥ 2 , where u0 is the ground-truth IC and u∗0 is the optimizer iterate at the reported checkpoint. Solver color and marker encode the backend; line style encodes the optimizer. Adam stagnates above 40% normalized IC error across all solvers, while L-BFGS converges below 6% in fewer gradient evaluations. Adding a solenoidal projection after each L-BFGS update (L-BFGS+proj) reduces the final IC error to below 0.5% and keeps the recovered IC divergence at the physical level of the true IC throughout. The projection benefit is specific to the second-order optimizer, and the asymmetry has a structural cause. L-BFGS builds its Hessian approximation from secant pairs 𝑠𝑘 = 𝑥𝑘+1 − 𝑥𝑘 ; without projection these steps carry large irrotational components orthogonal to the solenoidal manifold, corrupting the curvature model, whereas projection confines 𝑠𝑘 to the correct tangent space and recovers the superlinear convergence behavior associated

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

Drag reduction

Optimised profile

Profile history XLB

1.0 0.8 40

PhiFlow

0.6 30

Drag reduction (%)

0.4 20 PICT

0.2 10 0.0

0 0

200

400

0.25

0.50

0.75

Iteration Initial

XLB

PhiFlow

250

500

Iteration

PICT

Figure 20 | 2D drag minimization at Re = 20. XLB and PICT converge to ∼60% drag reduction; PhiFlow stalls at ∼35%.

with projected Newton-type methods [7]. Adam’s per-coordinate second-moment scaling accumulates variance estimates from both solenoidal and irrotational gradient directions and has no cross-coordinate curvature model to exploit the constraint geometry [29], so projecting the iterate does not repair the miscalibrated scaling; Figure 22 confirms this. The bottom row shows 𝑢𝑥 at the middle 𝑧-slice for a representative solver (L-BFGS+proj/PhiFlow). Topology optimization. Figure 23 summarizes the structural topology optimization task ([0, 2] ×

[0, 1] × [0, 1] cantilever beam, 16 × 2 × 8 mesh, SIMP with 𝑝 = 3, Adam). All solvers converge to the same final compliance (≈0.001): for smooth, well-conditioned physics, gradient quality does not limit solution quality. The optimized density fields (bottom row) are visually identical across solvers, recovering the classic cantilever-beam truss topology. We also evaluated L-BFGS and MMA [49] on this task. After 100 quasi-Newton steps, L-BFGS final compliance ranged from 0.001 (TopOpt.jl) to 0.008 (JAX-FEM), well above the Adam solution. L-BFGS fails for two reinforcing reasons: the SIMP penalization (𝑝 = 3) makes the objective globally non-convex, so the positive-curvature condition 𝑦𝑘⊤𝑠𝑘 > 0 required for valid BFGS updates is frequently violated and updates are skipped [18]; and box constraints 𝜌 ∈ [0, 1] cause rapid active-set changes that invalidate the limited curvature history. MMA, by contrast, converges to compliance ≈0.00099 in only 40 iterations, matching Adam’s solution in under 2% of its iteration budget, because it builds a strictly convex separable approximation at each step that sidesteps the SIMP non-convexity and handles box constraints and the volume inequality natively. MMA is therefore the recommended optimizer for this task class, consistent with standard practice in the topology optimization community. Conductivity recovery. Figure 24 shows convergence for the thermal conductivity inversion

task. The ground-truth conductivity field has two Gaussian peaks; the solver observes the steady-state temperature distribution under a known heat source and recovers the conductivity by minimizing the ℓ2 error between simulated and observed temperature. Adam and L-BFGS are compared across solvers that support differentiation w.r.t. the conductivity field. FEniCS and Firedrake form a distinct cluster: with Adam their final error is modestly higher (∼ 1,600 vs. ∼ 1,300), and with L-BFGS the gap widens sharply (∼ 8.5×104 vs. ∼ 1,100). The cause is Í a gradient inconsistency in the adjoint-based solvers: the forward pass reports 𝑖 (𝑇𝑖 − 𝑇𝑖∗ ) 2

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

IC recovery error

IC divergence

Optimization loss

100 10−3

10−1 10−1

10−5 10−7

10−2

10−9 10

−3

102

Gradient evaluations

IC recovered

IC true

101

100

102

Gradient evaluations

Final state

0.45

Final true

0.23

0.00

-0.23

-0.45

Adam

L-BFGS+proj

INS.jl

PICT

L-BFGS

PhiFlow

XLB

Warp-NS

Exponax

Figure 21 | 3D NS IC recovery overview for Adam, L-BFGS, and L-BFGS+proj (Adam+proj omitted; see Figure 22). Top: normalized IC error, IC divergence, and optimization loss vs. gradient evaluations across all solvers (color = solver, line style = optimizer). The dashed line marks the true IC divergence level. Bottom: 𝑢𝑥 slice at convergence (L-BFGS+proj / PhiFlow): recovered IC, true IC, recovered final state, and true final state.

IC recovery error

IC divergence

Optimization loss 0.03

0.15

0.8

0.02

0.10 0.6

0.01

0.05

0.00 0

200

400

Gradient evaluations

200

400

Gradient evaluations

Adam

PhiFlow

XLB

Adam+proj

INS.jl

PICT

200

400

Gradient evaluations Warp-NS

Exponax

Figure 22 | Adam vs. Adam+proj on the 3D NS IC recovery task (linear scale). Adding a solenoidal projection step after each Adam update leaves IC error, IC divergence, and optimization loss essentially unchanged across all solvers, confirming that the projection benefit observed in Figure 21 is specific to L-BFGS.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026 Final fields, fine mesh (32×4×16, MMA) JAX-FEM

Convergence, coarse mesh (16×2×8)

TopOpt.jl

FEniCS

Firedrake

Compliance

0.005

Adam MMA

0.004

JAX-FEM TopOpt.jl

0.003

FEniCS Firedrake

0.002 0.001 100

101

102

103

Front (XZ)

Top (XY)

Right (YZ)

Volume fraction

Iteration 0.54 0.52 0.50

100

101

102

103

Iteration Fixed (u = 0)

Load

Figure 23 | Structural topology optimization overview. Left: final optimized density fields from a finer-mesh (32×4×16) MMA run, both as 3D voxel renderings (one per solver) and projected along each principal axis. Right: compliance and volume-fraction vs. iteration for Adam and MMA (line style = optimizer, color = solver) on the coarser 16×2×8 mesh used throughout the rest of the comparison; MMA converges in ∼40 iterations, Adam requires ∼2500. The two columns therefore use different mesh resolutions, as flagged in the panel headers: the finer run sharpens the recovered topology visually, while the coarser run keeps optimizer convergence directly comparable to the other tasks.

∫ (a nodal sum), while the adjoint differentiates (𝑇 − 𝑇 ∗ ) 2 dΩ (an area-weighted integral) and applies a scalar correction factor 𝑛 nodes /|Ω| to bridge the two. This correction is only approximate, so the returned gradient is not the exact derivative of the reported objective. First-order methods tolerate this mismatch because they use only gradient direction; L-BFGS accumulates gradient differences to estimate curvature, so an inconsistent gradient corrupts the Hessian approximation and causes convergence to a suboptimal point.

C. Solver documentation This appendix documents all solvers evaluated for inclusion in Mosaic, covering implementation details, wrapping approach, and known limitations. Solvers included in the benchmark appear in Table 2. Excluded solvers and the reasons for their exclusion are listed in Section C.2. C.1 Included solvers JAX-CFD [30]. JAX-native incompressible flow solver on a staggered MAC grid with finite-

difference advection and spectral (FFT) pressure projection. Gradients via JAX source-transformation AD. The spectral pressure solve requires periodic BCs in all spatial directions, so JAX-CFD is included in periodic benchmarks (TGV, multimode agreement) but excluded from the cylinderwake experiment.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

Thermal conductivity recovery

Identification error

107 106 105 104 103 101

100

102

103

Gradient evaluations

Conductivity

Adam — recovered profiles 0.8 0.6 0.4 0.2 0.0 0

100

120

100

120

FEniCS

TorchFEM

Node index

Conductivity

L-BFGS — recovered profiles 0.8 0.6 0.4 0.2 0.0 0

Node index Truth

Firedrake

JAX-FEM

Figure 24 | Thermal conductivity recovery comparing Adam and L-BFGS optimizers. Top: identification error vs. gradient evaluations; color encodes solver, linestyle encodes optimizer. Bottom: final recovered conductivity profiles (all solvers) vs. ground truth (dashed).

PhiFlow [24]. Semi-Lagrangian advection with pressure projection. Supports PyTorch, JAX,

and TensorFlow backends, enabling gradient computation through the same simulation code across AD frameworks.

INS.jl [1]. IncompressibleNavierStokes.jl: Julia finite-difference pressure-projection solver. Dif-

ferentiates through the time loop via Zygote.jl reverse-mode AD. CPU only.

XLB [4]. JAX-native lattice Boltzmann solver supporting D2Q9 (2D) and D3Q27 (3D) stencils

on GPU. Gradients via source-transformation AD through the collision-streaming loop.

PICT [17]. GPU-accelerated differentiable incompressible Navier-Stokes solver built on PyTorch

with custom CUDA kernels implementing the PISO algorithm. Supports reverse-mode AD through the full time loop and handles multi-block curvilinear grids. The Mosaic Tesseract

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

exposes periodic, lid-driven cavity, inflow/channel (with advective outflow), and cylinder-wake modes via an 8-block ring topology. Warp-NS [37]. Custom incompressible NS solver built with NVIDIA Warp CUDA kernels, im-

plementing the IPCS projection scheme with an FFT-based Poisson solve (following a reference spectral example; a non-periodic solver was not implemented). Gradients via the wp.Tape kernel-level VJP. The Warp framework provides differentiable kernels but no ready-made incompressible flow solver, so the FD stencils, pressure Poisson solve, and time integration were implemented from scratch using Warp primitives. Warp-NS represents a best-effort implementation by the Mosaic team and illustrates the integration cost when a kernel toolkit provides no built-in solver.

Exponax [31]. Spectral incompressible Navier-Stokes solver supporting both 2D (streamfunctionvorticity) and 3D (Leray-projected velocity) formulations, integrated with an exponential timedifferencing Runge-Kutta (ETDRK) scheme. Incompressibility is enforced to machine precision by construction. Gradients via JAX source-transformation AD. Used in Mosaic for the 3D periodic domain (F3). FEniCS [16]. Finite-element solver using P1 elements for thermal and structural problems.

dolfin-adjoint automates the discrete adjoint by replaying the forward tape.

Firedrake [48]. Mirrors the FEniCS P1/CG1 formulation for structural and thermal problems.

Differentiates via firedrake-adjoint, providing an independent tape-based adjoint implementation for cross-validation. JAX-FEM [54]. Solves heat conduction and linear elasticity with trilinear HEX8 finite elements

in JAX. Gradients via AD through the assembled system.

TopOpt.jl [25]. SIMP topology optimization for linear elasticity with HEX8 elements in Julia,

using analytical adjoint sensitivities.

torch-fem [38]. PyTorch finite-element solver for heat conduction with linear HEX8 elements.

Gradients via PyTorch autograd through the assembled system. GPU-accelerated via PyTorch sparse operations.

OpenFOAM [52] (reference). Runs the icoFoam incompressible PISO solver as a forward-only

reference baseline. No reverse-mode AD available.

deal.II [3, 5] (reference). Solves thermal and structural problems with Q1 elements using the

industry-grade C++ finite-element library. No native reverse-mode AD is available; used as a forward-only reference baseline for S and H, analogously to OpenFOAM for the fluid domains. C.2 Excluded solvers

We document solvers that were evaluated for inclusion but ultimately excluded, along with the specific reason in each case. Practitioners looking for a particular solver can check whether it was considered and why it was left out. WaterLily.jl [53]. Forward-mode AD only, via ForwardDiff.jl dual numbers. No Zygote reverse-

mode support is documented. A VJP can be emulated by contracting 𝑁 forward-mode JVPs, but the cost scales as 𝑂 (𝑁 𝑑 · 𝑇fwd ), matching central finite differences. Excluded from gradient benchmarks. Forward accuracy is reported where applicable. JAX-Fluids [9]. Targets compressible flow. Recovering the incompressible limit at low Mach

number (Ma ≈ 0.01) requires acoustic sub-stepping whose iteration count depends on Ma, grid size, and time step simultaneously. Because jax.lax.scan requires a statically known loop count, a fully differentiable low-Mach wrapper cannot be built around the published library without reimplementing its time integrator. Excluded from both forward and gradient benchmarks.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

Commercial solvers (ANSYS, COMSOL, Abaqus). Excluded for reproducibility: closed-source

licenses prevent redistribution of solver binaries or independent verification of gradient implementations. FEniCS-NS [16]. The FEniCS incompressible NS solver (CG1/CG1 elements, IPCS time-stepping)

was evaluated for fluid benchmarks but excluded for two independent reasons. Without SUPG/GLS advection stabilization, stability requires Peℎ = 𝑈 ℎ/(2𝜈) ≲ 1; at 𝜈 < 0.01 on the 𝑁 = 64 benchmark grid, relative errors reach 31–212 while other solvers stay in 0.02–0.16. Independently, wall-clock scaling is empirically 𝑂 (𝑁 2.45 ) in serial; the VJP at 𝑁 = 64 takes 1248 s, and the forward pass alone for the drag optimization task (𝑁 = 32, 400 steps) requires approximately 25,000 s, exceeding the 1200 s HTTP watchdog by a factor of 40. FEniCS remains included for structural and heat-transfer domains (H, S), where neither constraint applies. C.3 Observed solver limitations

The following limitations were identified during benchmarking and resulted in task-specific exclusions. Solvers excluded globally are documented in Section C.2. PICT: viscosity not differentiable in PISOtorch_diff. PISOtorch_diff tracks autograd through velocity fields and boundary conditions, but treats viscosity as a static scalar: passing a requires_grad=True tensor as the viscosity argument produces a result with no grad_fn (confirmed by direct inspection of the autograd graph). Differentiation w.r.t. 𝜈 would require adding viscosity as an explicit differentiable parameter inside the PISOtorch_diff C++/CUDA kernels; this is outside the scope of the Mosaic wrapper. JAX-CFD: spectral pressure solve requires periodic boundary conditions. JAX-CFD uses an FFT-

based pressure Poisson solver that requires doubly-periodic boundary conditions by construction. The 2D cylinder drag optimization domain is a channel with inlet, outlet, and obstacle boundaries, none of which are periodic. Volume penalization masks velocity inside the solid region but does not remove the periodicity requirement from the pressure solve. JAX-CFD is structurally excluded from the drag optimization task. Warp-NS: no non-periodic pressure solver available within the time budget. The Warp-NS imple-

mentation was built following a reference example that uses an FFT-based spectral Poisson solver, which inherits the same doubly-periodic BC requirement. Unlike JAX-CFD, this is not a structural limitation of the Warp framework: a non-periodic iterative Poisson solver could in principle be implemented using Warp primitives. However, no suitable example was available, and implementing one from scratch was outside the time budget of this work; the drag optimization was therefore not attempted. INS.jl: viscosity gradient not available. Zygote.jl cannot differentiate through the diffusion term in IncompressibleNavierStokes.jl because its ChainRulesCore rule for the diffusion operator returns NoTangent() for the viscosity argument. Gradients w.r.t. 𝜈 are therefore not reported

for INS.jl. Fixing this requires either upstreaming a correct rrule into INS.jl or implementing a custom adjoint that accounts for the diffusion term.

INS.jl: Brinkman penalization incompatible with spectral pressure solve. INS.jl (Incompressible-

NavierStokes.jl) [1] supports periodic, Dirichlet, symmetric, and advective outlet boundary conditions, but not immersed boundary (IBM) or cut-cell methods. Applying Brinkman volume penalization to represent the cylinder obstacle introduces a velocity discontinuity at the solid boundary after each spectral LU pressure projection step; the Poisson solve amplifies this discontinuity, and the velocity field diverges to NaN at all tested resolutions (𝑁 ≥ 16). A correct approach requires incorporating the penalization term directly into the pressure Poisson system [2], which is outside the scope of the current solver implementation. INS.jl is excluded from the drag optimization task.

Mosaic: A Benchmark Suite for Differentiable Physics Solvers — Rehmann et al., 2026

XLB: intrinsic 𝑂 (Ma2 ) compressibility error. The D2Q9 BGK lattice Boltzmann method [45]

recovers the incompressible Navier-Stokes equations only to 𝑂 (Ma2 ) via Chapman-Enskog expansion. At fixed Δ𝑡, the Mach number Ma = 𝑢 Δ𝑡/Δ𝑥 grows with grid refinement, so the compressibility error does not decrease with spatial resolution for incompressible problems. At 𝑁 = 128, Δ𝑡 = 0.01, Ma ≈ 0.2, giving an 𝑂 (0.04) error floor independent of solver correctness [4]. XLB forward accuracy results at fine grids should be interpreted with this limit in mind. XLB: BGK collision instability at low viscosity. The BGK relaxation time 𝜏 = 𝜈/𝑐𝑠2 Δ𝑡 + 0.5 ap-

proaches the stability boundary 𝜏 → 0.5 as 𝜈 → 0 [45]. The XLB tesseract automatically selects the KBC (entropic) collision operator [4] when 𝜔 > 1.8 (𝜏 < 0.556), which is unconditionally entropy-stable. Results at 𝜈 ∈ {10−4, 5 × 10−4 } are pending re-validation under the KBC operator and should be interpreted with caution.

Record · ID 319696 · SHA-256 5d54f952b32d6cf7
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.