Fix qk_clip_utils import

Author: hsuan-lun-chiang

src/maxtext/trainers/pre_train/train.py

@@ -62,7 +62,7 @@
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
-from MaxText.utils import qk_clip_utils
+from maxtext.utils import qk_clip_utils
 from maxtext.utils import sharding
 from maxtext.utils import train_utils
 from maxtext.utils.vocabulary_tiling import vocab_tiling_linen_loss

tests/unit/qk_clip_test.py

@@ -27,7 +27,7 @@
 from maxtext.common.gcloud_stub import is_decoupled
 from maxtext.layers import attention_mla
 from maxtext.utils import maxtext_utils
-from MaxText.utils.qk_clip_utils import apply_qk_clip, calculate_max_logit_metric
+from maxtext.utils.qk_clip_utils import apply_qk_clip, calculate_max_logit_metric
 
 from MaxText import pyconfig
 from tests.utils.test_helpers import get_test_config_path