flax
Fix ValueError when `nnx.jit` is used with `nnx.custom_vjp`
#5045
Merged

Loading