[mosaic_gpu] Warmup before measuring the running time in profiler.measure

PiperOrigin-RevId: 700650380

Author: superbobry

jax/experimental/mosaic/gpu/profiler.py

@@ -104,6 +104,7 @@ def run(*args, **kwargs):
       raise ValueError("Can only measure functions with at least one output")
     return outs, _event_elapsed(start_event, end_event)
 
+  jax.block_until_ready(run(*args, **kwargs))  # Warmup.
   outs, elapsed = run(*args, **kwargs)
   return outs, float(elapsed)