Fix incorrect argument in JAX flash attention. (#21014)

hertschuh · rtg0795 · commit 54bdd1916c07 · 2025-03-26T22:21:33.000Z
The mask is named `array` in `NumpyMask`.
diff --git a/keras/src/backend/jax/nn.py b/keras/src/backend/jax/nn.py
@@ -1124,7 +1124,7 @@ def wrap_flash_attention(
         )
 
     if custom_mask is not None:
-        mask = splash_attention_mask.NumpyMask(mask=custom_mask)
+        mask = splash_attention_mask.NumpyMask(array=custom_mask)
 
     else:
         mask = splash_attention_mask.CausalMask(

Original file line number	Diff line number	Diff line change
`@@ -1124,7 +1124,7 @@ def wrap_flash_attention(`
`1124`	`1124`	`)`
`1125`	`1125`
`1126`	`1126`	`if custom_mask is not None:`
`1127`		`- mask = splash_attention_mask.NumpyMask(mask=custom_mask)`
	`1127`	`+ mask = splash_attention_mask.NumpyMask(array=custom_mask)`
`1128`	`1128`
`1129`	`1129`	`else:`
`1130`	`1130`	`mask = splash_attention_mask.CausalMask(`