jax-fem icon indicating copy to clipboard operation
jax-fem copied to clipboard

Published 20 hours ago • verified publisher icon deepmodeling

PETSc solver not working with JAX transforms

Open SNMS95 opened this issue 9 months ago • 0 comments

  • 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

SNMS95 avatar Oct 03 '23 08:10 SNMS95