class_docstr: OmniMixer model.
  
  Attributes:
    num_classes: Number of output classes.
    patch_size: Patch size of the stem.
    hidden_size: Size of the hidden state of the output of model's stem.
    num_layers: Number of layers.
    channels_mlp_dim: hidden dimension of the channel mixing MLP.
    sequence_mlp_dim: hidden dimension of the token (sequence) mixing MLP.
    omnimixer: Configurations of the omnimixer (omnidirectional mixer).
    dropout_rate: Dropout rate.
    stochastic_depth: overall stochastic depth rate.
docstr: OmniMixer model.
docstr_tok: ["OmniMixer","model","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      *,
      train: bool,
      debug: Optional[bool] = False,
  ) -> jnp.ndarray:
      """OmniMixer model."""
      x = nn.Conv(
          self.hidden_size,
          self.patch_size,
          strides=self.patch_size,
          padding="VALID",
          name="embedding",
      )(x)
      n, h, w, c = x.shape
      x = jnp.reshape(x, [n, h * w, c])
  
      x = OmniMixerEncoder(
          num_layers=self.num_layers,
          channels_mlp_dim=self.channels_mlp_dim,
          sequence_mlp_dim=self.sequence_mlp_dim,
          omnimixer=self.omnimixer,
          dropout_rate=self.dropout_rate,
          stochastic_depth=self.stochastic_depth,
          name="omnimixer_encoder",
      )(x, train=train)
  
      # Use global average pooling for our classifier, dim (1,) or (1,2).
      x = jnp.mean(x, axis=list(range(1, x.ndim - 1)))
  
      x = nn_layers.IdentityLayer(name="pre_logits")(x)
      x = nn.Dense(
          self.num_classes,
          kernel_init=nn.initializers.zeros,
          name="output_projection",
      )(x)
      return x
  
imports: from typing import Any, Optional, Sequence
  from absl import logging
  import flax.linen
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import mixer
  from scenic.projects.baselines import vit
  from scenic.projects.fast_vit import model_utils
  from scenic.projects.omninet import model_utils
label_desc: Use image node_id 1 for calling the OmniMixer obj's underlying member method code with example usage: obj.__call__(x) and returns: x