Incorrect gradients for 3D Helmholtz simulations using checkpointing

[astanziola]

There are some edge cases for which Helmholtz simulations in 3D give incorrect gradients when checkpoint=True is used in the helmholtz_solver for FourierSeries fields.

This is likely to be a jax issue, since results were correct up to jax 0.3.20.
It has been raised upstream: google/jax#14302

For the moment, if memory requirements allow, setting checkpoint=False should get around the problem.

May make sense to add a warning when checkpoint=True and setting its default value to False, especially if this is going to take long to be solved from the jax side.

[astanziola]

Worth checking if this is still an issue with lineax.

[astanziola]

This seems to be fixed now. Added some regression tests to catch it if it happens again