Show HN: Finite Element Analysis with JAX
3 ichihara_naruki 2 8/7/2025, 10:33:22 AM github.com ↗
FEAX is a compact, high-performance finite element analysis engine built on JAX. It provides an API for solving partial differential equations with automatic differentiation, JIT compilation, and GPU acceleration.
For anything non-local (electromagnetics being my old research area for e.g.), for performance reasons you almost always have to do some sort of accelerated computation, because otherwise no interesting problem is tractable. In order to do that, you end up coupling some other library or bespoke code to the FE solver to use BEM method, Fast Multipole Method, etc. and these then break all the assumptions about differentiability and what allows you to solve adjoint problems. But you basically have no choice - it's either this or nothing.
At the same time, I started developing this library with a future in mind—where large-memory GPUs are affordable and ubiquitous. JAX has dramatically lowered the barrier to GPU computing, and with FEAX, users can instantly generate GPU-compatible FEA code.
Over the past few days, I’ve been testing it on mid-scale problems with vectorized batch boundary conditions, solving hundreds of linear systems in parallel. These are use cases where traditional FEM frameworks often fall short.
I hope that features like this can help accelerate interdisciplinary research at the intersection of physics simulation and AI.