class_docstr: OmniNet model.
  
  Attributes:
      num_classes: number of classes.
      mlp_dim: Dimension of the MLP on top of attention block.
      num_layers: Number of layers.
      attention_configs: Configurations passed to the self-attention.
      attention_fn: Self-attention function used in the model.
      patches: Configuration of the patches extracted in the stem of the model.
      hidden_size: Size of the hidden dimension on the stem of the model.
      omninet: Configurations of the omninet (omnidirectional attention).
      representation_size: Size of the final representation.
      dropout_rate: Dropout rate.
      attention_dropout_rate: Dropout rate for attention heads.
      classifier:  Type of the classifier.
docstr: OmniNet model.
docstr_tok: ["OmniNet","model","."]
code: def __call__(
      self,
      inputs: jnp.ndarray,
      *,
      train: bool,
      debug: Optional[bool] = False,
  ) -> jnp.ndarray:
      """OmniNet model."""
      patch_height, patch_width = self.patches.size
      patch_stride_height, patch_stride_width = self.patches.get(
          "strides", self.patches.size
      )
      x = nn.Conv(
          self.hidden_size,
          (patch_height, patch_width),
          strides=(patch_stride_height, patch_stride_width),
          padding="VALID",
          name="embedding",
      )(inputs)
  
      # Flatten the input.
      bs, h, w, c = x.shape
      x = jnp.reshape(x, [bs, h * w, c])
  
      # If we want to add a class token, add it here.
      if self.classifier == "token":
          cls = self.param(
              "cls", nn.initializers.zeros, (1, 1, c), x.dtype
          )
          cls = jnp.tile(cls, [bs, 1, 1])
          x = jnp.concatenate([cls, x], axis=1)
  
      x = OmnidirectionalEncoder1D(
          mlp_dim=self.mlp_dim,
          num_layers=self.num_layers,
          attention_configs=self.attention_configs,
          attention_fn=self.attention_fn,
          omninet=self.omninet,
          dropout_rate=self.dropout_rate,
          attention_dropout_rate=self.attention_dropout_rate,
          stochastic_droplayer_rate=self.stochastic_droplayer_rate,
          name="Transformer",
      )(x, train=train)
  
      if self.classifier in ("token", "0"):
          x = x[:, 0]
      elif self.classifier in ("gap", "gmp", "gsp"):
          fn = {"gap": jnp.mean, "gmp": jnp.max, "gsp": jnp.sum}[
              self.classifier
          ]
          x = fn(x, axis=1)
  
          x = jnp.mean(x, axis=list(range(1, x.ndim - 1)))
  
      if self.representation_size is not None:
          x = nn.Dense(self.representation_size, name="pre_logits")(x)
          x = nn.tanh(x)
      else:
          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 OmniNet obj's underlying member method code with example usage: obj.__call__(inputs) and returns: x