fix lint

Cjkkkk · Cjkkkk · commit 8c4d6d6903a3 · 2025-02-03T06:09:05.000Z
diff --git a/jax/_src/cudnn/fused_attention_stablehlo.py b/jax/_src/cudnn/fused_attention_stablehlo.py
@@ -1789,7 +1789,7 @@ def dot_product_attention(
           f"but got: bias={bias}, mask={mask}, q_seqlen={q_seqlen}, kv_seqlen={kv_seqlen}"
       )
     check_fp8_params(fp8_params)
-    check_layout(query, key, value, bias, q_seqlen, kv_seqlen, layout)
+    check_layout(query, key, value, bias, q_seqlen, kv_seqlen, q_offsets, kv_offsets, layout)
     output, amax_s, amax_o = _dot_product_attention_fp8(
         query, key, value, fp8_params,
         scale, mask_type == MaskType.CAUSAL, layout.value, cudnn_version

Original file line number	Diff line number	Diff line change
`@@ -1789,7 +1789,7 @@ def dot_product_attention(`
`1789`	`1789`	`f"but got: bias={bias}, mask={mask}, q_seqlen={q_seqlen}, kv_seqlen={kv_seqlen}"`
`1790`	`1790`	`)`
`1791`	`1791`	`check_fp8_params(fp8_params)`
`1792`		`- check_layout(query, key, value, bias, q_seqlen, kv_seqlen, layout)`
	`1792`	`+ check_layout(query, key, value, bias, q_seqlen, kv_seqlen, q_offsets, kv_offsets, layout)`
`1793`	`1793`	`output, amax_s, amax_o = _dot_product_attention_fp8(`
`1794`	`1794`	`query, key, value, fp8_params,`
`1795`	`1795`	`scale, mask_type == MaskType.CAUSAL, layout.value, cudnn_version`