Finite-JAX: Differentiable Finite Solvers

• Finite-JAX is a JAX-based framework offering explicit and implicit finite-difference solvers for PDEs like the Burgers and Navier–Stokes equations.
• It integrates both external-solver coupling (e.g., Firedrake–JAX) and pure-JAX implementations, ensuring end-to-end differentiability for sensitivity analysis and optimization.
• The framework leverages JIT compilation and automatic differentiation to enhance performance across CPUs, GPUs, and TPUs while supporting real-world PDE-constrained applications.

Finite-JAX, also written FiniteJAX, denotes a JAX-based finite-difference solver framework in which explicit and implicit solvers are implemented for the one-dimensional viscous Burgers equation, and the same label has also been used for a JAX-implemented two-dimensional incompressible Navier–Stokes channel-flow solver. Related work further treats the connection of mature finite-discretization software to JAX automatic differentiation as a concrete realization of the same underlying idea: finite discretization solvers exposed as differentiable JAX computations (Potluri et al., 18 Jan 2026, Rodriguez et al., 29 Sep 2025, Yashchuk, 2023).

1. Terminology and scope

In the shock-wave-physics usage, Finite-JAX is a finite-difference solver framework implemented in JAX, not just a single solver. It includes explicit and implicit finite-difference solvers for the one-dimensional viscous Burgers equation, with upwind discretization of the convective term, central differences for diffusion, and explicit and implicit time integration. The framework is built on JAX, which “provides automatic differentiation, just-in-time (JIT) compilation, and device-agnostic execution across CPUs, GPUs, and TPUs” (Potluri et al., 18 Jan 2026).

A second usage appears in wall-bounded-flow verification, where Finite-JAX is a two-dimensional incompressible Navier–Stokes solver implemented in JAX using finite differences. There it is applied to laminar, pressure-driven channel flow, with cross-model verification against the analytical Hagen–Poiseuille solution and benchmarking on CPU, GPU, and TPU backends (Rodriguez et al., 29 Sep 2025).

A broader interpretation is made explicit in the Firedrake–JAX work, which is described as “essentially a concrete realization of what you are calling ‘Finite-JAX’”: a pattern in which a mature finite element PDE solver is connected to JAX so that the PDE solve becomes a differentiable JAX primitive with custom forward and reverse rules (Yashchuk, 2023). This suggests that the term is best understood not only as the name of specific finite-difference codes, but also as a differentiable-programming paradigm for finite discretizations in JAX.

2. Core computational principles

A central formulation for the broader Finite-JAX idea is the implicit PDE map

$F(u,m)=0,$

where $u$ is the discretized PDE solution and $m$ collects parameters such as material coefficients, control fields, or neural-network weights. In the Firedrake–JAX interface, the PDE residual is expressed symbolically in UFL, and differentiation is performed not by tracing every Newton iteration or Krylov step, but by solving the associated tangent-linear and adjoint equations. For forward mode, the Jacobian–vector product is obtained from

$$\frac{\partial F(u,m)}{\partial u}\,\dot u_v = -\frac{\partial F(u,m)}{\partial m}\cdot v,$$

and for reverse mode the vector–Jacobian product is obtained from

$$\frac{\partial F(u,m)}{\partial u}^{*}\lambda_v = v,\qquad \frac{du}{dm}^{*}v = -\frac{\partial F(u,m)}{\partial m}^{*}\lambda_v.$$

This construction makes the Firedrake solve appear as a JAX primitive with custom JVP and VJP rules, while bypassing low-level solver internals (Yashchuk, 2023).

A distinct but related pattern appears in fully JAX-native solvers. JAX-FVM states that every operation—mesh connectivity, flux evaluation, slope limiting, and time integration—is written entirely in JAX, so that the whole solver is just-in-time compiled, vectorised, and end-to-end differentiable through automatic differentiation, and runs transparently on CPU or GPU (Romémont, 8 Jul 2026). The literature therefore distinguishes between two architectural regimes: external-solver coupling with custom differentiation rules, and pure-JAX discretizations in which AD traverses the entire numerical pipeline directly.

3. Finite-difference realizations under the Finite-JAX name

In the Burgers-equation implementation, the governing PDE is

$$\frac{\partial u}{\partial t} + u\frac{\partial u}{\partial x} = \nu\frac{\partial^2 u}{\partial x^2},$$

used as a simplified analogue of the Navier–Stokes equations for shock formation, propagation, and dissipation. The convective derivative is discretized by a backward difference, corresponding to upwinding for positive velocity,

$$\frac{\partial u}{\partial x}\bigg|_{x_i}\approx \frac{u_i^n-u_{i-1}^n}{\Delta x},$$

while the diffusive term uses the standard second-order central stencil,

$$\frac{\partial^2 u}{\partial x^2}\bigg|_{x_i}\approx \frac{u_{i+1}^n-2u_i^n+u_{i-1}^n}{\Delta x^2}.$$

Time integration is performed with explicit forward Euler and implicit BTCS, and the code also implements RK4 and a Crank–Nicolson-like implicit variant. The explicit update is

$$u_i^{n+1}=u_i^n-\Delta t\,u_i^n\frac{u_i^n-u_{i-1}^n}{\Delta x} +\nu\frac{\Delta t}{\Delta x^2}(u_{i+1}^n-2u_i^n+u_{i-1}^n).$$

The implicit BTCS formulation leads to a nonlinear tridiagonal system solved by Newton–Raphson; an iterative Gauss–Seidel variant stops when the maximum residual is below $10^{-6}$. Stability of the explicit method is enforced by the convective and diffusive CFL bounds

$u$0

For $u$1, $u$2, and $u$3, reported wall times and errors include: explicit CPU (NumPy) 0.92 s with $u$4, explicit CPU (JAX) 0.88 s with $u$5, explicit GPU (JAX L4) 0.17 s with $u$6, implicit GPU (JAX A100) 0.12 s with $u$7, and implicit TPU (v5e-1) 0.11 s with $u$8 (Potluri et al., 18 Jan 2026).

In the channel-flow study, Finite-JAX solves the two-dimensional incompressible Navier–Stokes equations with a constant streamwise body force $u$9,

$m$0

in a rectangular channel with no-slip walls at $m$1 and $m$2, periodicity in $m$3, $m$4, $m$5, $m$6, $m$7, and $m$8. The numerical method uses second-order central differences, explicit first-order Euler time stepping, and a Chorin projection with a pressure Poisson equation. The simulation is stopped when

$m$9

The analytical comparison target is the plane Poiseuille profile

$$\frac{\partial F(u,m)}{\partial u}\,\dot u_v = -\frac{\partial F(u,m)}{\partial m}\cdot v,$$0

with $$\frac{\partial F(u,m)}{\partial u}\,\dot u_v = -\frac{\partial F(u,m)}{\partial m}\cdot v,$$1. Reported verification errors are $$\frac{\partial F(u,m)}{\partial u}\,\dot u_v = -\frac{\partial F(u,m)}{\partial m}\cdot v,$$2, max norm $$\frac{\partial F(u,m)}{\partial u}\,\dot u_v = -\frac{\partial F(u,m)}{\partial m}\cdot v,$$3, and min norm $$\frac{\partial F(u,m)}{\partial u}\,\dot u_v = -\frac{\partial F(u,m)}{\partial m}\cdot v,$$4. Hardware times for the Euler discretization are approximately 6.215 s on NVIDIA A100, 4.728 s on NVIDIA L4, 4.394 s on NVIDIA T4, 0.935 s on TPU v6e-1, 1.133 s on TPU v5e-1, and 1.585 s on CPU (Rodriguez et al., 29 Sep 2025).

4. Finite elements, finite volumes, and the broader differentiable-discretization ecosystem

The Firedrake–JAX interface illustrates how finite elements enter the same conceptual space. PDEs are written in UFL, solved by Firedrake and PETSc, and wrapped as differentiable JAX functions. The paper demonstrates Poisson optimal control, using SciPy’s L-BFGS-B with JAX gradients and convergence in 38 iterations, and coefficient inversion where a finite-element parameterization with 981 parameters converges in 31 iterations while a neural-network parameterization with 47 parameters converges in 89 iterations; the neural representation gives better reconstruction of the true $$\frac{\partial F(u,m)}{\partial u}\,\dot u_v = -\frac{\partial F(u,m)}{\partial m}\cdot v,$$5 and less overfitting to noise (Yashchuk, 2023).

The available literature also suggests an extension of the Finite-JAX idea from structured finite differences to unstructured finite volumes and differentiable structural FEM. JAX-FVM is an open-source, fully differentiable finite volume method for the two-dimensional compressible Euler and Navier–Stokes equations on unstructured triangular meshes, built around an entropy-conservative Tadmor/Ismail-Roe two-point flux, entropy-variable Rusanov or Roe dissipation, second-order MUSCL reconstruction, and explicit and matrix-free implicit time integrators whose Jacobian actions are obtained by AD (Romémont, 8 Jul 2026). JAX-SSO is a differentiable finite element analysis solver for structural optimization and seamless integration with neural networks, with examples in shape optimization, thickness optimization, simultaneous shape and topology optimization, and physics-informed neural networks (Wu, 2024). JAX-FEM is a differentiable GPU-accelerated three-dimensional FEM library; in a 3D tensile loading problem with 7.7 million degrees of freedom, JAX-FEM with GPU achieves around $$\frac{\partial F(u,m)}{\partial u}\,\dot u_v = -\frac{\partial F(u,m)}{\partial m}\cdot v,$$6 acceleration compared to a commercial FEM code depending on platform (Xue et al., 2022).

5. Applications and scientific-machine-learning workflows

The principal application domains attached to Finite-JAX-style methods are PDE-constrained optimization, inverse problems, parameter estimation, and hybrid neural-network–PDE models. In the Firedrake–JAX setting, differentiable PDE solves are used for optimal control, coefficient-field inversion, and neural-network parameterizations inside PDE constraints (Yashchuk, 2023). In the Burgers and channel-flow studies, the same differentiable-programming structure is positioned toward sensitivity analysis, gradient-based parameter estimation, and integration into physics-informed neural networks or hybrid CFD–ML training loops (Potluri et al., 18 Jan 2026, Rodriguez et al., 29 Sep 2025).

A transport-equation example sharpens the scientific-machine-learning angle. The JAX-based finite-volume solver for population balance equations is fully differentiable, achieves up to $$\frac{\partial F(u,m)}{\partial u}\,\dot u_v = -\frac{\partial F(u,m)}{\partial m}\cdot v,$$7 relative acceleration in PBE simulations versus NumPy, and shows that differentiability can be $$\frac{\partial F(u,m)}{\partial u}\,\dot u_v = -\frac{\partial F(u,m)}{\partial m}\cdot v,$$8 faster for optimizing larger models than conventional approaches. The solver is explicitly presented as a foundation for hybrid models in which learnable components replace empirical growth laws or act as in-the-loop error correctors (Alsubeihi et al., 2024).

A more hybrid construction appears in JAX-DIPS, where neural networks are trained not by AD-based differential residuals, but by finite discretization residuals evaluated on implicit Cartesian cells centered on collocation points. This “neural bootstrapping method” is described as competitive in terms of memory and training speed with other PINN-type frameworks and is applied to elliptic interface problems with discontinuities in three spatial dimensions (Mistani et al., 2022). This suggests a second application route for Finite-JAX: not only differentiable solvers, but also differentiable training objectives derived directly from finite discretizations.

6. Limitations, misconceptions, and future directions

A recurrent distinction in the literature is between external-solver coupling and native JAX discretization. The Firedrake bridge is explicit that the solver is not written in JAX; it is external, based on Firedrake and PETSc, and the PDE part is not an XLA-compiled JAX kernel. Differentiation relies on Firedrake’s UFL and dolfin-adjoint/tlm_adjoint, and the approach is tailored to variational FEM formulations. The same source notes ongoing work on tighter integration, just-in-time compilation, GPU computing, and better interoperability between JAX’s parallelism and Firedrake’s MPI parallelism (Yashchuk, 2023).

A second common misconception is that Finite-JAX already denotes a mature, standardized CFD platform. The named finite-difference studies remain narrowly scoped: one-dimensional viscous Burgers shock dynamics with explicit and implicit schemes, and two-dimensional laminar channel flow with explicit Euler time stepping and a projection method (Potluri et al., 18 Jan 2026, Rodriguez et al., 29 Sep 2025). Their stated future directions include extension from one-dimensional Burgers to multi-dimensional Navier–Stokes, more complex boundary conditions and geometries, higher-order or shock-capturing schemes such as ENO/WENO and flux limiters, three-dimensional wall-bounded flows, improved Poisson solvers, and hybrid ML–CFD strategies (Potluri et al., 18 Jan 2026, Rodriguez et al., 29 Sep 2025).

Recent pure-JAX finite-volume work identifies additional frontier problems: JAX-FVM is presently restricted to two-dimensional triangular meshes, with three-dimensional extension, turbulence modeling, and general VTK mesh input listed as planned developments (Romémont, 8 Jul 2026). Taken together, the literature suggests that Finite-JAX is best understood as an emerging differentiable finite-discretization paradigm inside the JAX ecosystem, rather than as a single canonical solver.