Include phase corrections in jax

Fixes #29 . This MR should not fix the inherent issues in the phase correction, but should rather match the results of JAX and numpy in the phase corrections.

This is now done, and can be merged as the next step towards v0.8