class_docstr: Tests the default trainer on single device setup.
class_docstr_tok: ["Tests","the","default","trainer","on","single","device","setup","."]
docstr: Test training for two epochs on MNIST with a small model.
docstr_tok: ["Test","training","for","two","epochs","on","MNIST","with","a","small","model","."]
code: def test_trainer(self):
      """Test training for two epochs on MNIST with a small model."""
  
      rng = jax.random.PRNGKey(0)
      np.random.seed(0)
      config = ml_collections.ConfigDict(
          {
              "dataset_name": "mnist",
              "data_dtype_str": "float32",
              "rng_seed": 0,
              "lr_configs": {
                  "learning_rate_schedule": "compound",
                  "factors": "constant * cosine_decay",
                  "steps_per_cycle": 100,
                  "base_learning_rate": 0.1,
              },
              "hid_sizes": [20, 10],
              "model_dtype_str": "float32",
              "optimizer": "momentum",
              "optimizer_configs": {"momentum": 0.9},
              "batch_size": 128,
              "eval_batch_size": 64,
              "l2_decay_factor": 0.0005,
              "max_grad_norm": None,
              "label_smoothing": None,
              "write_summary": None,  # no summary writing
              "checkpoint": False,  # no checkpointing
              "debug_eval": False,
              "debug_train": False,
              "xprof": False,
          }
      )
  
      with tfds.testing.mock_data(num_examples=1024):
          model_cls = models.get_model_cls(
              "fully_connected_classification"
          )
          dataset_builder = datasets.get_dataset("mnist")
  
          dataset = dataset_builder(
              batch_size=config.batch_size,
              eval_batch_size=config.eval_batch_size,
              num_shards=jax.local_device_count(),
              dtype_str=config.data_dtype_str,
          )
  
          config.num_training_steps = 100
          config.log_eval_steps = 50
          config.num_training_epochs = None
          _, train_summary, eval_summary = classification_trainer.train(
              rng=rng,
              config=config,
              model_cls=model_cls,
              dataset=dataset,
              workdir=self.test_dir,
              writer=metric_writers.LoggingWriter(),
          )
  
      self.assertGreaterEqual(train_summary["accuracy"], 0.0)
      self.assertLess(train_summary["loss"], 5.0)
      self.assertGreaterEqual(eval_summary["accuracy"], 0.0)
      self.assertLess(eval_summary["loss"], 5.0)
  
imports: import functools
  import shutil
  import tempfile
  from absl.testing import absltest
  from clu import metric_writers
  import flax
  from flax import jax_utils
  import flax.linen
  import jax.numpy
  import jax.random
  import ml_collections
  import numpy
  from scenic.dataset_lib import datasets
  from scenic.model_lib import models
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import multilabel_classification_model
  from scenic.train_lib import classification_trainer
  from scenic.train_lib import lr_schedules
  from scenic.train_lib import optimizers
  from scenic.train_lib import train_utils
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 7 for calling the ClassificationTrainerTest obj's underlying member method code with example usage: obj.test_trainer() without return types