Add back the import of jtu in flash_attention.py

This was erroneously removed in de3191f.

Author: andportnoy

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

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