Changing the return type

parambole · web-flow · commit 70beed7c83f9 · 2025-09-30T12:24:06.000-07:00
diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py
@@ -395,7 +395,7 @@ def _cudnn_flash_attention(query: Array, key: Array, value: Array, heads: int, m
       check_rep=False,
   )
   def wrap_flash_attention(query, key, value):
-    return jax.vmap(dpa_layer)(query, key, value, mask=None)
+    return dpa_layer(query, key, value, mask=None)
 
   out = wrap_flash_attention(query, key, value)
   return _reshape_data_from_cudnn_flash(out)

Original file line number	Diff line number	Diff line change
`@@ -395,7 +395,7 @@ def _cudnn_flash_attention(query: Array, key: Array, value: Array, heads: int, m`
`395`	`395`	`check_rep=False,`
`396`	`396`	`)`
`397`	`397`	`def wrap_flash_attention(query, key, value):`
`398`		`- return jax.vmap(dpa_layer)(query, key, value, mask=None)`
	`398`	`+ return dpa_layer(query, key, value, mask=None)`
`399`	`399`
`400`	`400`	`out = wrap_flash_attention(query, key, value)`
`401`	`401`	`return _reshape_data_from_cudnn_flash(out)`