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 "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({})
      optimizer = optimizers.get_optimizer(config).create(
          initial_params
      )
      init_train_state = jax_utils.replicate(
          train_utils.TrainState(
              global_step=0,
              optimizer=optimizer,
              model_state=init_model_state,
              rng=jax.random.PRNGKey(0),
          )
      )
      return FakeFlaxModel(), init_train_state
  
imports: import functools
  import shutil
  import os
  from absl.testing import absltest
  from absl.testing import parameterized
  import flax
  from flax import jax_utils
  import flax.linen
  import jax.numpy
  import jax.random
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import multilabel_classification_model
  from scenic.projects.vivit import evaluation_lib
  from scenic.projects.vivit import train_utils
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import train_utils
  import tensorflow
label_desc: Use image node_id 3 for calling the ViViTClassificationTrainerTest 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