PETSc solver not working with JAX transforms

deepmodeling / jax-fem

Differentiable Finite Element Method with JAX

GNU General Public License v3.0

290 stars 45 forks source link

Open SNMS95 opened 1 year ago

SNMS95 commented 1 year ago

Currently, the solver [when using petsc] cannot be combined with any of Jax's transforms
This prevents vmapping and higher order derivatives
Solution:
- Wrap petsc with Lineax or
- Wrap it a pure-callback + custom-jvp