Jax-FEM: A differentiable GPU-accelerated 3D FEM solver (github.com)

Paper describing this and what it’s for here:

https://arxiv.org/pdf/2212.00964 (PDF)

We emphasize the following three features that differentiate JAX-FEM from other FEM libraries:

1. Efficient solution to forward PDE with GPU acceleration;

2. Differentiable simulation for automatic inverse design;

3. Seamless integration with machine learning