class_docstr: Tests for modules in mixer.py.
class_docstr_tok: ["Tests","for","modules","in","mixer.py","."]
docstr: Tests MixerBlock.
docstr_tok: ["Tests","MixerBlock","."]
code: def test_mixer_block(
      self, dropout_rate, stochastic_depth, layer_scale=None
  ):
      """Tests MixerBlock."""
      rng = random.PRNGKey(0)
      x = jnp.ones((4, 16, 32))
      mixer_block = functools.partial(
          mixer.MixerBlock,
          channels_mlp_dim=32,
          sequence_mlp_dim=32,
          dropout_rate=dropout_rate,
          stochastic_depth=stochastic_depth,
          layer_scale=layer_scale,
      )
      mixer_block_vars = mixer_block().init(rng, x, deterministic=True)
      y = mixer_block().apply(mixer_block_vars, x, deterministic=True)
      # Test outputs shape.
      self.assertEqual(y.shape, x.shape)
  
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 1 for calling the MixerTest obj's underlying member method code with example usage: obj.test_mixer_block(dropout_rate, stochastic_depth, layer_scale) without return types