flax

fix partial eval for custom_jvp and omnistaging changes

#393

Merged