Add back the import of jtu in flash_attention.py

This was erroneously removed in de3191f.
jax-ml · Oct 22, 2024 · 6378984 · 6378984
1 parent 2b7b074
commit 6378984
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/jax/experimental/mosaic/gpu/examples/flash_attention.py b/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