Merge branch 'main' of github.com:AI-Hypercomputer/maxtext into shuningjin-gmm-clean

Author: shuningjin

tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_mt.sh

@@ -39,6 +39,7 @@ python3 -m maxtext.checkpoint_conversion.to_maxtext \
     scan_layers=false \
     hardware=cpu skip_jax_distributed_system=True \
     checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
+    --lazy_load_tensors=False \
     --eager_load_method='transformers'
 
 UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/unscanned/${run_id}/0/items
@@ -52,6 +53,7 @@ python3 -m maxtext.checkpoint_conversion.to_maxtext \
     scan_layers=true \
     hardware=cpu skip_jax_distributed_system=True \
     checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
+    --lazy_load_tensors=False \
     --eager_load_method='transformers'
 
 SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/scanned/${run_id}/0/items

tests/post_training/unit/train_distill_test.py

@@ -845,6 +845,7 @@ def __call__(self, x):
     with self.assertRaises(AssertionError, msg="Weights should have updated on the second pass."):
       np.testing.assert_allclose(student.linear.kernel.get_value(), initial_weights)
 
+  @pytest.mark.skip(reason="Hangs indefinitely on synthetic datasets")
   @mock.patch("clu.metric_writers.create_default_writer")
   @mock.patch("maxtext.trainers.post_train.distillation.train_distill.tokenizer.build_tokenizer")
   def test_train_save_and_resume(self, mock_build_tokenizer, mock_writer):