Chunk and local attention fix for auto-regressive generation

Author: gagika

MaxText/layers/attentions.py

@@ -874,8 +874,12 @@ def cudnn_jax_flash_attention(
       decoder_segment_ids: Array | None,
       model_mode: str = MODEL_MODE_TRAIN,
   ) -> tuple[Array, Array]:
+<<<<<<< HEAD
     """CUDNN Flash Attention with JAX SDPA API.
     """
+=======
+    """CUDNN Flash Attention with JAX SDPA API."""
+>>>>>>> 2560bdc6 (Chunk and local attention fix for auto-regressive generation)
     # These imports are only meant to work in a GPU build.
     # pylint: disable=import-outside-toplevel
     from jax._src.cudnn.fused_attention_stablehlo import (