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

Signed-off-by: Kunjan Patel <kunjanp@google.com>

Author: coolkp

src/maxdiffusion/tests/wan_transformer_test.py

@@ -197,33 +197,15 @@ def test_wan_block(self):
  assert dummy_output.shape == dummy_hidden_states.shape
 
  def test_wan_attention(self):
- pyconfig.initialize(
- [
- None,
- os.path.join(THIS_DIR, "..", "configs", "base_wan_14b.yml"),
- ],
- unittest=True,
- )
- config = pyconfig.config
+ # pyconfig.initialize(
+ #  [
+ #  None,
+ #  os.path.join(THIS_DIR, "..", "configs", "base_wan_14b.yml"),
+ #  ],
+ #  unittest=True,
+ # )
+ # config = pyconfig.config
 
- batch_size = 1
- channels = 16
- frames = 21
- height = 90
- width = 160
- hidden_states_shape = (batch_size, frames, height, width, channels)
- dummy_hidden_states = jnp.ones(hidden_states_shape)
- wan_rot_embed = WanRotaryPosEmbed(attention_head_dim=128, patch_size=[1, 2, 2], max_seq_len=1024)
- dummy_rotary_emb = wan_rot_embed(dummy_hidden_states)
-
- key = jax.random.key(0)
- rngs = nnx.Rngs(key)
- devices_array = create_device_mesh(config)
-
- mesh_axes = ['data', 'fsdp', 'tensor']
- mesh = Mesh(devices_array, mesh_axes)
- batch_size = 1
- query_dim = 5120
  for attention_kernel in ["flash", "tokamax_flash"]:
  pyconfig.initialize(
  [
@@ -233,6 +215,22 @@ def test_wan_attention(self):
  ]
  )
  config = pyconfig.config
+ batch_size = 1
+ channels = 16
+ frames = 21
+ height = 90
+ width = 160
+ hidden_states_shape = (batch_size, frames, height, width, channels)
+ dummy_hidden_states = jnp.ones(hidden_states_shape)
+ wan_rot_embed = WanRotaryPosEmbed(attention_head_dim=128, patch_size=[1, 2, 2], max_seq_len=1024)
+ dummy_rotary_emb = wan_rot_embed(dummy_hidden_states)
+
+ key = jax.random.key(0)
+ rngs = nnx.Rngs(key)
+ devices_array = create_device_mesh(config)
+ mesh = Mesh(devices_array, config.mesh_axes)
+ batch_size = 1
+ query_dim = 5120
  with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
  flash_block_sizes = get_flash_block_sizes(config)
  attention = FlaxWanAttention(