Copybara import of the project:

--
81f4467d19e3ba22927aa1bb1597877229adc7b8 by A9isha <mazumdera@google.com>:

Add GRPO demo script, refactor create_nnx_model

COPYBARA_INTEGRATE_REVIEW=#2258 from AI-Hypercomputer:anisha-grpo-script 81f4467d19e3ba22927aa1bb1597877229adc7b8
PiperOrigin-RevId: 801670635

Author: A9isha

MaxText/__init__.py

@@ -25,12 +25,12 @@
 from collections.abc import Sequence
 
 from MaxText import maxtext_utils
-from MaxText import train_utils
+from MaxText import model_creation_utils
 from MaxText import pyconfig
 from MaxText.layers import models
 from MaxText import dpo_utils
 from jax.sharding import Mesh
-from MaxText.train_utils import from_config
+from MaxText.model_creation_utils import from_config
 
 Transformer = models.Transformer
 transformer_as_linen = models.transformer_as_linen