mlsanityData to 7 Oct 2026

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

arXiv 2212.00964 ↗First posted 2 Dec 2022cs.MS cs.CEcited by 40

Not on the map. Why?

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.

Key facts. Cited by 40 papers in the map (data through 7 Oct 2026). First posted on arXiv 2 Dec 2022.

Builds on

No reference list yet.

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2023: 3; 2024: 6; 2025: 16; 2026: 15.

40 papers.

Sort:

By area:

  1. 2024Neural Material Adaptor for Visual Grounding of Intrinsic Dynamicscited by 30
  2. 2023XLB: A differentiable massively parallel lattice Boltzmann library in Pythoncited by 11
  3. 2026Autonomous Agents Coordinating Distributed Discovery Through Emergent Artifact Exchangecited by 11
  4. 2025Bohrium + SciMaster: Building the Infrastructure and Ecosystem for Agentic Science at Scalecited by 8
  5. 2023Generative Hyperelasticity with Physics-Informed Probabilistic Diffusion Fieldscited by 7
  6. 2024JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flowscited by 7
  7. 2025MASIV: Toward Material-Agnostic System Identification from Videoscited by 7
  8. 2025Polyconvex Physics-Augmented Neural Network Constitutive Models in Principal Stretchescited by 5
  9. 2025History-Aware Neural Operator: Robust Data-Driven Constitutive Modeling of Path-Dependent Materialscited by 5
  10. 2025JAX-MPM: A Learning-Augmented Differentiable Meshfree Framework for GPU-Accelerated Lagrangian Simulation and Geophysical Inverse Modelingcited by 5
  11. 2025PhysGen: Physically Grounded 3D Shape Generation for Industrial Designcited by 5
  12. 2023Neural-Integrated Meshfree (NIM) Method: A differentiable programming-based hybrid solver for computational mechanicscited by 4
28 more
  1. 2024Differentiable programming across the PDE and Machine Learning barriercited by 4
  2. 2025Scientific Machine Learning of Flow Resistance Using Universal Shallow Water Equations with Differentiable Programmingcited by 3
  3. 2025Physics-Informed DeepONet Coupled with FEM for Convective Transport in Porous Media with Sharp Gaussian Sourcescited by 3
  4. 2026Learning, Solving and Optimizing PDEs with TensorGalerkin: an efficient high-performance Galerkin assembly algorithmcited by 3
  5. 2026Towards Rapid Constitutive Model Discovery from Multi-Modal Data: Physics Augmented Finite Element Model Updating (paFEMU)cited by 3
  6. 2024Gradient-free neural topology optimization: Towards effective fracture-resistant designscited by 2
  7. 2024Differentiable Neural-Integrated Meshfree Method for Forward and Inverse Modeling of Finite Strain Hyperelasticitycited by 2
  8. 2025Physics-informed Reduced Order Modeling of Time-dependent PDEs via Differentiable Solverscited by 2
  9. 2025COMMET: orders-of-magnitude speed-up in finite element method via batch-vectorized neural constitutive updatescited by 2
  10. 2025Differentiable Inverse Modeling with Physics-Constrained Latent Diffusion for Heterogeneous Subsurface Parameter Fieldscited by 2
  11. 2024Learning Nonlinear Finite Element Solution Operators using Multilayer Perceptrons and Energy Minimizationcited by 1
  12. 2026FEM-Informed Hypergraph Neural Networks for Efficient Elastoplasticitycited by 1
  13. 2026AutoMS: Multi-Agent Evolutionary Search for Cross-Physics Inverse Microstructure Designcited by 1
  14. 2026KATOsuper: Surrogate-accelerated neural topology optimization with sensitivity-consistent Fourier neural operatorscited by 1
  15. 2025Fully data-driven inverse hyperelasticity with hyper-network neural ODE fieldscited by 0
  16. 2025MicroLad: 2D-to-3D Microstructure Reconstruction and Generation via Latent Diffusion and Score Distillationcited by 0
  17. 2025Estimating Parameter Fields in Multi-Physics PDEs from Scarce Measurementscited by 0
  18. 2025Case study of a differentiable heterogeneous multiphysics solver for a nuclear fusion applicationcited by 0
  19. 2025GLUE: Coordinating Pre-Trained Generative Models for System-Level Designcited by 0
  20. 2026Guided Diffusion by Optimized Loss Functions on Relaxed Parameters for Inverse Material Designcited by 0
  21. 2026Material-agnostic temperature field prediction for metal additive manufacturing via a parametric PINN frameworkcited by 0
  22. 2026Differentiable Multiphysics Co-Optimization via Implicit Neural Representations: A Transient Hamburger-Cooking Benchmarkcited by 0
  23. 2026Inverse Design of Metainterfaces for Static Friction Control: Beyond the Hertzian Limitcited by 0
  24. 2026Finite Element-Based Material Learning via Automatic Differentiation: Learning constitutive neural network models from full-field deformation datacited by 0
  25. 2026A Physics-Informed Neural Network Framework for Elastodynamic Wave Propagation in Bimaterial Systemscited by 0
  26. 2026Learning Implicit Constitutive Laws for Dynamic 3D Gaussian Splatting from Monocular Videoscited by 0
  27. 2026Benchmarking Optimizers to Solve Inverse Problems with Differentiable Physics Simulatorscited by 0
  28. 2026CraftSPH: A high-accuracy and composable differentiable SPH solver implemented in PyTorchcited by 0