JAX-FEM: differentiating a finite element solver for inverse design

Automatic differentiation plus the finite element method

JAX-FEM is a differentiable finite element package based on JAX, and the README states the concept in one line: it is Automatic Differentiation plus Finite Element Method.

That combination is what separates it from an ordinary finite element code. In a conventional solver, the simulation is a function you evaluate and the sensitivities you need for an optimisation are something you derive by hand, which for a nonlinear multi-physics problem is a substantial piece of work and a common source of errors. In JAX-FEM, the simulation is written as a JAX program, so the gradients come out of the same code path.

The stated capability is differentiable programming for solving inverse and design problems without deriving sensitivities by hand, with topology optimization and optimal thermal control named as the examples. The other examples in the README show the range of physics: a thermal profile in direct energy deposition, linear static analysis of a bracket, crystal plasticity with a grain structure and a stress field, Stokes flow with velocity and pressure, and a topology optimization animated with a differentiable simulation.

The audience is a research group or an engineering team that wants to optimise against a physics simulation, rather than a team that just wants answers for a fixed structure. The package is built for automatic differentiation and GPU execution, and the package description calls it a GPU accelerated finite element analysis package in JAX.

What the package covers and what it leaves out

The key features list in the README is specific enough to check a requirement against.

On discretisation, it supports 2D quadrilateral and triangle elements and 3D hexahedron and tetrahedron elements, and it handles first and second order elements. On boundary conditions, Dirichlet, Neumann and Robin are all supported. On physics, it covers linear and nonlinear analysis including the heat equation, linear elasticity, hyperelasticity, and plasticity in both the macro and crystal flavours. Multi-physics problems are supported, and there is an integration with PETSc for solver options.

Two things stand out for an engineering reader. The plasticity entry is not a generic mention, it names macro and crystal plasticity, which are different constitutive models with different state variables, so the package reaches into a real part of the mechanics toolbox rather than only linear elasticity. And PETSc matters more than it might look, because the linear solver behind a finite element assembly is often where the run time actually goes, and having a choice of solver options is what lets you tune a model for the machine you have.

The update note in the README, dated December 11, 2023, is worth reading because it describes the multi-physics work in detail: multiple variables can be solved monolithically, weak forms are defined through volume and surface integrals, and body force, mass kernel and Laplace kernel are treated in a unified way through the volume integral while Neumann and Robin boundary conditions are treated in a unified way through the surface integral.

The project also points at JAX-FEM Express, an LLM agent built for the package, hosted as an application, with a video walkthrough.

The design and optimisation cases the package targets

The value of a differentiable solver is that the design loop becomes ordinary code. If the forward solve is a JAX function, then a gradient of an objective with respect to a design variable is another JAX call, and an optimiser step is a third.

The README names the applications it has in mind. Topology optimization with differentiable simulation is the headline case, and the animated example in the README is exactly that loop in motion. Optimal thermal control is the second, and the direct energy deposition thermal profile example shows the kind of problem where a design decision, here how to deposit material, has a thermal consequence you want to minimise.

The bracket example and the crystal plasticity example are there to show that the same machinery handles a structural and a material problem, not only a thermal one. The Stokes flow example adds a fluid problem, with velocity and pressure fields as the outputs, which is a case where the adjoint or reverse-mode formulation has to work across a coupled field.

So the scope is not a single physics. It is a framework in which several physics are written the same way, and the differentiation works for all of them.

The papers the README asks you to cite make the same point from a research angle, with the 2023 Computer Physics Communications article describing JAX-FEM as a differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science.

Installing the package and finding the guide

The README does not carry installation commands. It says that for installation and the user guide you should visit the documentation at deepmodeling.github.io/jax-fem, and that is the reference for both.

What the repository shows is the shape of the project. The package itself lives in jax_fem/, the example applications in applications/, the tests in tests/, and the documentation sources in docs/. There is an environment.yml for a Conda environment and a pyproject.toml for the Python package.

The build configuration is short and modern:

[build-system]
requires = ["hatchling==1.26.3", "hatch-vcs"]
build-backend = "hatchling.build"

hatch-vcs means the version comes from the version control tags rather than being written by hand, and the pinned hatchling version means the build itself is reproducible.

The project metadata says the package is jax-fem, that it requires Python 3.10 or later, and that it is described as a GPU accelerated finite element analysis package in JAX. The authors listed are Tianju Xue and Weipeng Xu.

One inconsistency in the metadata is worth flagging, and it concerns licensing rather than code. The classifiers in pyproject.toml include the OSI approved BSD License line, while the README says the project is licensed under the GNU General Public License v3 and points at the LICENSE file for details. The pyproject also sets the license field to that file. The README is the authoritative statement, and it adds that for commercial use you should contact Tianju Xue.

Where JAX-FEM is the wrong tool

Three cases point elsewhere.

The first is a conventional analysis workflow. If you want a stress result for one bracket under one load case, a mature finite element package with validated element formulations, a post-processor and a qualification record is a better fit. A differentiable solver earns its complexity when you are going to differentiate, and the README is explicit that that is the motivation.

The second is a hardware constraint. The package is built for GPU execution, and the trade that buys is memory bandwidth and parallel throughput. A large nonlinear model with many degrees of freedom can be limited by device memory in a way that a CPU-oriented code with out-of-core solvers is not, and the README does not discuss fallback execution.

The third is a validation question, and it is the one to ask first. The examples in the README are demonstrations of capability, and a demonstration is not a qualification. Before this solves a structural problem that somebody's job depends on, you would want to confirm that the element formulation, the constitutive model and the boundary conditions match what your verification and validation process expects, and that the solutions converge to the reference result you already trust.

There is also a maintenance signal to weigh. The package version in the repository is 0.0.12 and the repository publishes no GitHub releases, so the version line in the source tree is the only reference, and the code is moving under a 0.x version where interfaces can still change. For a research group that pins an environment and follows the main branch, that is fine. For an engineering organisation that needs a supported release, it is a question to ask first.

Editorial conclusion

Adopt JAX-FEM if you need finite element simulation that participates in an optimisation loop, such as topology optimisation or inverse problems, and you are already working in JAX and can accept the Python 3.10 floor and the GPU memory model that comes with it. Do not adopt it as a general purpose structural solver for a conventional engineering workflow, and do not assume the solved problems are the ones your validation workflow covers. Verify three things first: how your constitutive model and element type map onto the package's supported set, since plasticity and hyperelasticity are specific formulations; which solver backend you will use, because the PETSc integration is an option rather than the default path; and the licensing question, because the project is GPL-3.0 with a commercial use route through the maintainer, while the package metadata carries a BSD classifier that disagrees. The documentation is the install guide, the last push was on 2026-10-09, and the package version in the repository is 0.0.12.

Frequently asked questions

Official sources

1. deepmodeling/jax-fem on GitHub
2. Issues
3. License: GPL-3.0
4. Project website
5. README