JAX related fixes

[nestordemeure]

• small performance improvement to some JAX kernels
• reintroduced datamovement tracker
• improved mergin of timing csv
• updated accel jax reset
• minor cleanups