remove common_types

Author: Cjkkkk

MaxText/layers/attentions.py

@@ -872,7 +872,7 @@ def cudnn_jax_flash_attention(
       key: Array,
       value: Array,
       decoder_segment_ids: Array | None,
-      model_mode: str = common_types.MODEL_MODE_TRAIN,
+      model_mode: str = MODEL_MODE_TRAIN,
   ) -> Array:
     """CUDNN Flash Attention with JAX SDPA API.
     """
@@ -885,7 +885,7 @@ def cudnn_jax_flash_attention(
 
     _, _, _, head_dim = query.shape  # pylint: disable=unused-variable
 
-    if model_mode == common_types.MODEL_MODE_AUTOREGRESSIVE:
+    if model_mode == MODEL_MODE_AUTOREGRESSIVE:
       lengths = jnp.sum(decoder_segment_ids, axis=-1)
 
       return dot_product_attention(