class_docstr: Tests for modules in mixer.py.
class_docstr_tok: ["Tests","for","modules","in","mixer.py","."]
docstr: Test forward pass of the mixer classification model.
docstr_tok: ["Test","forward","pass","of","the","mixer","classification","model","."]
code: def test_mixer_models(self):
      """Test forward pass of the mixer classification model."""
      rng = jax.random.PRNGKey(0)
      model = mixer.MixerMultiLabelClassificationModel(
          config=None,
          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 with mutated model_state.
      rng, dropout_rng = jax.random.split(rng)
      variables = {"params": init_params, **init_model_state}
      outputs, new_model_state = model.flax_model.apply(
          variables,
          xs,
          mutable=["batch_stats"],
          train=True,
          rngs={"dropout": dropout_rng},
          debug=False,
      )
      self.assertEqual(outputs.shape, (INPUT_SHAPE[0], NUM_OUTPUTS))
  
      # If it's a batch norm model check the batch stats changed.
      if init_model_state:
          bflat, _ = ravel_pytree(init_model_state)
          new_bflat, _ = ravel_pytree(new_model_state)
          self.assertFalse(jnp.array_equal(bflat, new_bflat))
  
      # Test batch_norm 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_mixer.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import flax
  from jax import random
  from jax.flatten_util import ravel_pytree
  import jax.numpy
  import jax.tree_util
  import numpy
  from scenic.projects.baselines import mixer
label_desc: Use image node_id 2 for calling the MixerTest obj's underlying member method code with example usage: obj.test_mixer_models() without return types