class_docstr: Tests the default trainer on single device setup.
class_docstr_tok: ["Tests","the","default","trainer","on","single","device","setup","."]
docstr: Generates the initial training state.
docstr_tok: ["Generates","the","initial","training","state","."]
code: def get_train_state(self, rng, fake_batch_logits):
      """Generates the initial training state."""
      config = ml_collections.ConfigDict(
          {
              "lr_configs": {
                  "base_learning_rate": 0.1,
              },
              "optimizer": "sgd",
          }
      )
  
      # define a fake model that always outputs the same logits: fake_batch_logits
      class FakeFlaxModel(nn.Module):
          """A fake flax model."""
  
          @nn.compact
          def __call__(self, x, train=False, debug=False):
              del x
              del train
              del debug
              # FakeFlaxModule always predicts class 2.
              return fake_batch_logits
  
      dummy_input = jnp.zeros((10, 10), jnp.float32)
      initial_params = (
          FakeFlaxModel()
          .init(rng, dummy_input)
          .get("params", flax.core.frozen_dict.FrozenDict({}))
      )
      init_model_state = flax.core.frozen_dict.FrozenDict({})
      lr_fn = lr_schedules.get_learning_rate_fn(config)
      optimizer_config = optimizers.get_optax_optimizer_config(config)
      tx = optimizers.get_optimizer(optimizer_config, lr_fn)
      opt_state = jax.jit(tx.init, backend="cpu")(initial_params)
      init_train_state = jax_utils.replicate(
          train_utils.TrainState(
              global_step=0,
              params=initial_params,
              tx=tx,
              opt_state=opt_state,
              model_state=init_model_state,
              rng=jax.random.PRNGKey(0),
          )
      )
      return FakeFlaxModel(), init_train_state
  
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 3 for calling the ClassificationTrainerTest obj's underlying member method code with example usage: obj.get_train_state(rng, fake_batch_logits) and returns: FakeFlaxModel,  init_train_state, fake_batch_logits