Test on attention type and automatically modify flash block sizes object when 'tokamax_flash' requested

coolkp · coolkp · commit a71eb74f44a3 · 2025-11-14T17:12:16.000Z
Signed-off-by: Kunjan Patel &lt;kunjanp@google.com&gt;
diff --git a/src/maxdiffusion/tests/wan_transformer_test.py b/src/maxdiffusion/tests/wan_transformer_test.py
@@ -234,7 +234,6 @@ def test_wan_attention(self):
       )
       config = pyconfig.config
       with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
-        config.attention = attention_kernel
         flash_block_sizes = get_flash_block_sizes(config)
         attention = FlaxWanAttention(
             rngs=rngs,

Original file line number	Diff line number	Diff line change
`@@ -234,7 +234,6 @@ def test_wan_attention(self):`
`234`	`234`	`)`
`235`	`235`	`config = pyconfig.config`
`236`	`236`	`with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):`
`237`		`- config.attention = attention_kernel`
`238`	`237`	`flash_block_sizes = get_flash_block_sizes(config)`
`239`	`238`	`attention = FlaxWanAttention(`
`240`	`239`	`rngs=rngs,`