class_docstr: Tests for modules in omninet model.py.
class_docstr_tok: ["Tests","for","modules","in","omninet","model.py","."]
docstr: Tests validity of output's shape of OmniMixerEncoder module.
docstr_tok: ["Tests","validity","of","output","'s","shape","of","OmniMixerEncoder","module","."]
code: def test_omnimixer_output_shape(self, skip_standard, partition, pool):
      """Tests validity of output's shape of OmniMixerEncoder module."""
      rng = random.PRNGKey(0)
      x = jnp.ones((4, 16, 32))
      omnimixer_configs = ml_collections.ConfigDict(
          {
              "skip_standard": skip_standard,
              "partition": partition,
              "pool": pool,
              "depth_mlp_dim": 16,
          }
      )
  
      omnimixer_encoder_def = functools.partial(
          model.OmniMixerEncoder,
          num_layers=4,
          channels_mlp_dim=32,
          sequence_mlp_dim=8,
          omnimixer=omnimixer_configs,
      )
      omnimixer_encoder_vars = omnimixer_encoder_def().init(
          rng, x, train=False
      )
      y = omnimixer_encoder_def().apply(
          omnimixer_encoder_vars, x, train=False
      )
      # Test outputs shape.
      self.assertEqual(y.shape, x.shape)
  
file_path: scenic/scenic/projects/omninet/tests/test_model.py
imports: import functools
  import itertools
  from absl.testing import absltest
  from absl.testing import parameterized
  from jax import random
  import jax.numpy
  import ml_collections
  from scenic.projects.omninet import model
label_desc: Use image node_id 1 for calling the OmniNetModelTest obj's underlying member method code with example usage: obj.test_omnimixer_output_shape(skip_standard, partition, pool) without return types