Merge conflict error

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

Author: coolkp

src/maxdiffusion/models/attention_flax.py

@@ -195,8 +195,14 @@ def _tpu_flash_attention(
         block_q_dkv=min(q_max_block_size, query.shape[2]),
         block_kv_dkv=min(kv_max_block_size, key.shape[2]),
         block_kv_dkv_compute=min(kv_max_block_size, query.shape[2]),
+<<<<<<< Updated upstream
         block_q_dq=min(q_max_block_size, query.shape[2]),
         block_kv_dq=min(kv_max_block_size, query.shape[2]),
+=======
+        block_q_dq=None if attention_kernel == "tokamax_flash" else min(q_max_block_size, query.shape[2]),
+        block_kv_dq=None if attention_kernel == "tokamax_flash" else min(kv_max_block_size, query.shape[2]),
+        use_fused_bwd_kernel=True if attention_kernel == "tokamax_flash" else False,
+>>>>>>> Stashed changes
     )
   num_fsdp_shards = mesh.shape["fsdp"]
   query = _reshape_data_for_flash(query, heads)