address lint and typecheck

Author: Cjkkkk

jax/_src/cudnn/fused_attention_stablehlo.py

@@ -369,7 +369,7 @@ def check_is_flash_attention(
           )
 
         if is_packed and cudnn_version < 90600:
-          raise NotImplementedError(f"Packed layout requires cudnn version >= 9.6.")
+          raise NotImplementedError("Packed layout requires cudnn version >= 9.6.")
 
 def check_cudnn_version():
   # check if cuDNN is installed

tests/fused_attention_stablehlo_test.py

@@ -616,6 +616,13 @@ def test_sdpa_large_head_size(self):
 
   @jtu.run_on_devices("cuda")
   def test_sdpa_packed_layout(self):
+    if jax.device_count() < 4:
+      self.skipTest("Requires more than 4 devices.")
+    try:
+      cudnn_version = check_cudnn_version()
+    except RuntimeError as e:
+      self.skipTest(str(e))
+      return
     if cudnn_version < 90600:
       self.skipTest("Requires >= cuDNN 9.6.0")
     k1, k2, k3, k4 = jax.random.split(jax.random.key(0), 4)