[Mosaic GPU] Add missing import.

PiperOrigin-RevId: 702629468

Author: dimitar-asenov

jax/experimental/mosaic/gpu/examples/matmul.py

@@ -20,6 +20,7 @@
 
 import jax
 from jax import random
+from jax._src import test_util as jtu  # noqa: F401
 from jax._src.interpreters import mlir
 from jax.experimental.mosaic.gpu import profiler
 from jax.experimental.mosaic.gpu import *  # noqa: F403