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.
Comments
physicsguy•2d ago
One of the biggest problems I've had with promises around this sort of framework approach for FE or FD or FVM is that in practice, they only can be used for the relatively narrow set of problems with local only phenomena. This is usually influenced by the biases of the fields the authors work in to be honest.
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.
ichihara_naruki•2d ago
I completely agree. For large-scale problems, specialized methods remain essential, and that’s exactly why reliable simulation tools will continue to play a critical role.
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.
physicsguy•2d ago
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.
ichihara_naruki•2d ago
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.