Fix for lower versions

Author: ae-foster

folx/experimental/pallas/attention/utils.py

@@ -7,7 +7,7 @@
 try:
     from jax.experimental.pallas import triton as plgpu
 except ImportError:
-    pass
+    from jax.experimental.pallas import gpu as plgpu
 
 
 from packaging.version import Version