class_docstr: Tests for modules in vit.py.
docstr: Test forward pass of the ViT classification model.
code: def test_vit_model(self, positional_embedding, use_default_config):
      """Test forward pass of the ViT classification model."""
      rng = jax.random.PRNGKey(0)
  
      # Config for model.
      if not use_default_config:
          config = ml_collections.ConfigDict(
              {
                  "model": dict(
                      num_heads=2,
                      num_layers=1,
                      representation_size=16,
                      mlp_dim=32,
                      dropout_rate=0.0,
                      attention_dropout_rate=0.0,
                      hidden_size=16,
                      patches={"size": (4, 4)},
                      positional_embedding=positional_embedding,
                      classifier="gap",
                      data_dtype_str="float32",
                  )
              }
          )
      else:
          config = None
      model = vit.ViTMultiLabelClassificationModel(
          config=config,
          dataset_meta_data={
              "num_classes": NUM_OUTPUTS,
              "target_is_onehot": False,
          },
      )
  
      xs = jnp.array(
          np.random.normal(loc=0.0, scale=10.0, size=INPUT_SHAPE)
      ).astype(jnp.float32)
  
      rng, init_rng = jax.random.split(rng)
      dummy_input = jnp.zeros(INPUT_SHAPE, jnp.float32)
      init_model_state, init_params = flax.core.pop(
          model.flax_model.init(
              init_rng, dummy_input, train=False, debug=False
          ),
          "params",
      )
  
      # Check that the forward pass works in train mode.
      rng, dropout_rng = jax.random.split(rng)
      variables = {"params": init_params, **init_model_state}
      outputs = model.flax_model.apply(
          variables,
          xs,
          train=True,
          rngs={"dropout": dropout_rng},
          debug=False,
      )
      self.assertEqual(outputs.shape, (INPUT_SHAPE[0], NUM_OUTPUTS))
  
      # Test model in inference mode.
      outputs = model.flax_model.apply(
          variables, xs, mutable=False, train=False, debug=False
      )
      self.assertEqual(outputs.shape, (INPUT_SHAPE[0], NUM_OUTPUTS))
  
file_path: scenic/scenic/projects/baselines/tests/test_vit.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import flax
  import jax
  import jax.numpy
  import jax.tree_util
  import ml_collections
  import numpy
  from scenic.projects.baselines import vit
label_desc: Use image node_id 1 for calling the ViTTest obj's underlying member method code with example usage: obj.test_vit_model(positional_embedding, use_default_config) without return types