[flax:examples:lm1b_nnx] Update example to work internally. #jax-fixit.

PiperOrigin-RevId: 840808821

Author: danielsuo

README.md

@@ -62,7 +62,7 @@ comes with everything you need to start your research, including:
 
 * **Utilities and patterns**: replicated training, serialization and checkpointing, metrics, prefetching on device.
 
-* **Educational examples**: [MNIST](https://flax.readthedocs.io/en/latest/mnist_tutorial.html), [Inference/sampling with the Gemma language model (transformer)](https://github.com/google/flax/tree/main/examples/gemma), [Transformer LM1B](https://github.com/google/flax/tree/main/examples/lm1b_nnx).
+* **Educational examples**: [MNIST](https://flax.readthedocs.io/en/latest/mnist_tutorial.html), [Inference/sampling with the Gemma language model (transformer)](https://github.com/google/flax/tree/main/examples/gemma).
 
 ## Quick install
 

examples/gemma/README.md

@@ -2,7 +2,7 @@
 ## Language modeling
 Trains Gemma model on the One Billion Word Benchmark (lm1b; Chelba *et al.*, 2013).
 
-This example is based on `lm1b_nnx` example script and similarly uses linear learning rate warmup and inverse square root learning rate schedule.
+This example is based on lm1b and similarly uses linear learning rate warmup and inverse square root learning rate schedule.
 
 
 ### Requirements

examples/gemma/input_pipeline_test.py

@@ -49,7 +49,7 @@ def _get_datasets(self):
     vocab_path = os.path.join(tempfile.mkdtemp(), 'sentencepiece_model')
 
     # Go two directories up to the root of the flax directory.
-    # "/path/to/flax/examples/lm1b_nnx/models_test.py" -> "/path/to/flax"
+    # "/path/to/flax/examples/gemma/input_pipeline_test.py" -> "/path/to/flax"
     flax_root_dir = pathlib.Path(__file__).absolute().parents[2]
     data_dir = str(flax_root_dir) + '/.tfds/metadata'  # pylint: disable=unused-variable