# JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science

arXiv: https://arxiv.org/abs/2212.00964 · HTML: https://mlsanity.com/paper/2212.00964/

## Key facts

- Cited by 41 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 2 Dec 2022.

## Built on by (most cited first, 30 of 41)

- [Neural Material Adaptor for Visual Grounding of Intrinsic Dynamics](https://mlsanity.com/paper/2410.08257/) (2024)
- [XLB: A differentiable massively parallel lattice Boltzmann library in Python](https://mlsanity.com/paper/2311.16080/) (2023)
- [Autonomous Agents Coordinating Distributed Discovery Through Emergent Artifact Exchange](https://mlsanity.com/paper/2603.14312/) (2026)
- [MASIV: Toward Material-Agnostic System Identification from Videos](https://mlsanity.com/paper/2508.01112/) (2025)
- [Bohrium + SciMaster: Building the Infrastructure and Ecosystem for Agentic Science at Scale](https://mlsanity.com/paper/2512.20469/) (2025)
- [Generative Hyperelasticity with Physics-Informed Probabilistic Diffusion Fields](https://mlsanity.com/paper/2310.03745/) (2023)
- [JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows](https://mlsanity.com/paper/2402.05193/) (2024)
- [PhysGen: Physically Grounded 3D Shape Generation for Industrial Design](https://mlsanity.com/paper/2512.00422/) (2025)
- [Polyconvex Physics-Augmented Neural Network Constitutive Models in Principal Stretches](https://mlsanity.com/paper/2503.00575/) (2025)
- [History-Aware Neural Operator: Robust Data-Driven Constitutive Modeling of Path-Dependent Materials](https://mlsanity.com/paper/2506.10352/) (2025)
- [JAX-MPM: A Learning-Augmented Differentiable Meshfree Framework for GPU-Accelerated Lagrangian Simulation and Geophysical Inverse Modeling](https://mlsanity.com/paper/2507.04192/) (2025)
- [Neural-Integrated Meshfree (NIM) Method: A differentiable programming-based hybrid solver for computational mechanics](https://mlsanity.com/paper/2311.12915/) (2023)
- [Differentiable programming across the PDE and Machine Learning barrier](https://mlsanity.com/paper/2409.06085/) (2024)
- [Scientific Machine Learning of Flow Resistance Using Universal Shallow Water Equations with Differentiable Programming](https://mlsanity.com/paper/2502.12396/) (2025)
- [Physics-Informed DeepONet Coupled with FEM for Convective Transport in Porous Media with Sharp Gaussian Sources](https://mlsanity.com/paper/2508.19847/) (2025)
- [Learning, Solving and Optimizing PDEs with TensorGalerkin: an efficient high-performance Galerkin assembly algorithm](https://mlsanity.com/paper/2602.05052/) (2026)
- [Towards Rapid Constitutive Model Discovery from Multi-Modal Data: Physics Augmented Finite Element Model Updating (paFEMU)](https://mlsanity.com/paper/2604.07746/) (2026)
- [Gradient-free neural topology optimization: Towards effective fracture-resistant designs](https://mlsanity.com/paper/2403.04937/) (2024)
- [Differentiable Neural-Integrated Meshfree Method for Forward and Inverse Modeling of Finite Strain Hyperelasticity](https://mlsanity.com/paper/2407.11183/) (2024)
- [Physics-informed Reduced Order Modeling of Time-dependent PDEs via Differentiable Solvers](https://mlsanity.com/paper/2505.14595/) (2025)
- [COMMET: orders-of-magnitude speed-up in finite element method via batch-vectorized neural constitutive updates](https://mlsanity.com/paper/2510.00884/) (2025)
- [Differentiable Inverse Modeling with Physics-Constrained Latent Diffusion for Heterogeneous Subsurface Parameter Fields](https://mlsanity.com/paper/2512.22421/) (2025)
- [Learning Nonlinear Finite Element Solution Operators using Multilayer Perceptrons and Energy Minimization](https://mlsanity.com/paper/2412.04596/) (2024)
- [FEM-Informed Hypergraph Neural Networks for Efficient Elastoplasticity](https://mlsanity.com/paper/2602.07364/) (2026)
- [AutoMS: Multi-Agent Evolutionary Search for Cross-Physics Inverse Microstructure Design](https://mlsanity.com/paper/2603.27195/) (2026)
- [KATOsuper: Surrogate-accelerated neural topology optimization with sensitivity-consistent Fourier neural operators](https://mlsanity.com/paper/2609.27216/) (2026)
- [Fully data-driven inverse hyperelasticity with hyper-network neural ODE fields](https://mlsanity.com/paper/2506.08146/) (2025)
- [MicroLad: 2D-to-3D Microstructure Reconstruction and Generation via Latent Diffusion and Score Distillation](https://mlsanity.com/paper/2508.20138/) (2025)
- [Estimating Parameter Fields in Multi-Physics PDEs from Scarce Measurements](https://mlsanity.com/paper/2509.00203/) (2025)
- [Case study of a differentiable heterogeneous multiphysics solver for a nuclear fusion application](https://mlsanity.com/paper/2511.13262/) (2025)

## Abstract (arXiv)

This paper introduces JAX-FEM, an open-source differentiable finite element method (FEM) library. Constructed on top of Google JAX, a rising machine learning library focusing on high-performance numerical computing, JAX-FEM is implemented with pure Python while scalable to efficiently solve problems with moderate to large sizes. For example, in a 3D tensile loading problem with 7.7 million degrees of freedom, JAX-FEM with GPU achieves around 10× acceleration compared to a commercial FEM code depending on platform. Beyond efficiently solving forward problems, JAX-FEM employs the automatic differentiation technique so that inverse problems are solved in a fully automatic manner without the need to manually derive sensitivities. Examples of 3D topology optimization of nonlinear materials are shown to achieve optimal compliance. Finally, JAX-FEM is an integrated platform for machine learning-aided computational mechanics. We show an example of data-driven multi-scale computations of a composite material where JAX-FEM provides an all-in-one solution from microscopic data generation and model training to macroscopic FE computations. The source code of the library and these examples are shared with the community to facilitate computational mechanics research.