Fix GM gpu timings
Replaces the generic jax_fori loop with cumsums for the WP and IWP
On Friday, 17th April, 9 am, the MPCDF GitLab instance will have a maintenance phase and will not be available for a short time.
Replaces the generic jax_fori loop with cumsums for the WP and IWP