Merge conflict error

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

Author: coolkp

src/maxdiffusion/models/attention_flax.py

@@ -195,14 +195,9 @@ 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)