  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$\times$
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.
