class_docstr: Omnidirectional Mixer Encoder.
  
  Attributes:
    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: The layer dropout rate (= stochastic depth).
  
  Returns:
    Output after OmniMixer block.
docstr: Applies OmniMixer encoder model on the inputs.
docstr_tok: ["Applies","OmniMixer","encoder","model","on","the","inputs","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      *,
      train: bool,
  ) -> jnp.ndarray:
      """Applies OmniMixer encoder model on the inputs."""
      partition_layers = []
      skip = (
          False  # First layer should be always a standard mixer layer.
      )
      for lyr in range(self.num_layers):
          if skip and self.omnimixer.skip_standard:
              logging.info("Skipping vanilla mixer at layer %d", lyr)
              skip = False
          else:
              p = (
                  lyr / max(self.num_layers - 1, 1)
              ) * self.stochastic_depth
              x = mixer.MixerBlock(
                  channels_mlp_dim=self.channels_mlp_dim,
                  sequence_mlp_dim=self.sequence_mlp_dim,
                  dropout_rate=self.dropout_rate,
                  stochastic_depth=p,
                  name=f"mixerblock_{lyr}",
              )(x, deterministic=not train)
              partition_layers.append(x)
          if len(partition_layers) == self.omnimixer.partition - 1:
              # Mixing in depth every N-1th layers.
              xp = jnp.stack(
                  partition_layers, 3
              )  # Shape: (bs, N, d, L).
              xp = attention_layers.MlpBlock(
                  mlp_dim=self.omnimixer.depth_mlp_dim,
                  dropout_rate=self.dropout_rate,
                  activation_fn=nn.gelu,
                  name=f"depth_mixing_{lyr}",
              )(xp, deterministic=not train)
              if self.omnimixer.pool == "last":
                  xp = xp[:, :, :, -1]  # Shape: (bs, N, d).
              elif self.omnimixer.pool == "max":
                  xp = jnp.max(xp, axis=3)  # Shape: (bs, N, d).
              else:
                  raise ValueError(
                      f"The pool type {self.omnimixer.pool} is not defined."
                  )
              x += xp
              partition_layers = []
              skip = True
      return nn.LayerNorm(name="encoder_norm")(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 OmniMixerEncoder obj's underlying member method code with example usage: obj.__call__(x) and returns: nn