class_docstr: Omnidirectional Encoder.
  
  Attributes:
    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.
    omninet: Configurations of the omninet (omnidirectional attention).
    dropout_rate:  Dropout rate.
    attention_dropout_rate: Dropout for attention heads.
docstr: Applies Transformer model on the inputs.
docstr_tok: ["Applies","Transformer","model","on","the","inputs","."]
code: def __call__(
      self,
      inputs: jnp.ndarray,
      train: bool,
  ) -> jnp.ndarray:
      """Applies Transformer model on the inputs."""
      assert inputs.ndim == 3  # Shape is `[batch, len, emb]`.
      x = vit.AddPositionEmbs(
          posemb_init=nn.initializers.normal(stddev=0.02),  # from BERT.
          name="posembed_input",
      )(inputs)
      x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not train)
  
      # Input encoder.
      partition_layers = []
      skip = False
      for lyr in range(self.num_layers):
          droplayer_p = (
              lyr / max(self.num_layers - 1, 1)
          ) * self.stochastic_droplayer_rate
          if skip and self.omninet.skip_standard:
              logging.info(
                  "Skipping vanilla transformer at layer %d", lyr
              )
              skip = False
          else:
              x = fast_vit_model_utils.Encoder1DBlock(  # pytype: disable=wrong-arg-types  # jax-ndarray
                  mlp_dim=self.mlp_dim,
                  attention_fn=self.attention_fn,
                  attention_configs=self.attention_configs,
                  dropout_rate=self.dropout_rate,
                  attention_dropout_rate=self.attention_dropout_rate,
                  droplayer_p=droplayer_p,
                  name=f"encoder_block_{lyr}",
              )(
                  x, inputs_kv=None, deterministic=not train
              )
              partition_layers.append(x)
  
          if len(partition_layers) == (self.omninet.partition - 1):
              # Every N-1th layer.
              #################### OmniNet ########################
              if self.omninet.query_type == "full":
                  if self.omninet.integrate == "concat":
                      xp = jnp.concatenate(partition_layers, 1)
                  elif self.omninet.integrate == "grid_stack":
                      xp = model_utils.grid_restack(partition_layers)
                  elif self.omninet.integrate == "factorized":
                      xp = jnp.stack(
                          partition_layers, 1
                      )  # Shape: (bs, L, N, d).
                  else:
                      raise ValueError(
                          f"The integrate type {self.omninet.integrate} is not defined."
                      )
  
                  def _pool(xp):
                      if self.omninet.pool == "last":
                          assert self.omninet.integrate == "concat"
                          return xp[:, -inputs.shape[1] :, :]
                      elif self.omninet.pool == "max":
                          assert self.omninet.integrate in [
                              "grid_stack",
                              "factorized",
                          ]
                          if self.omninet.integrate == "grid_stack":
                              return nn.max_pool(
                                  xp,
                                  (len(partition_layers),),
                                  strides=(len(partition_layers),),
                                  padding="VALID",
                              )
                          else:
                              return jnp.max(xp, axis=1)
                      else:
                          raise ValueError(
                              f"The pool type {self.omninet.pool} is not defined."
                          )
  
                  if self.omninet.pool_after_sa:
                      post_sa_fn = _pool
                  else:
                      post_sa_fn = None
  
                  if self.omninet.encoder.type == "1d":
                      xp = fast_vit_model_utils.Encoder1DBlock(  # pytype: disable=wrong-arg-types  # jax-ndarray
                          mlp_dim=self.omninet.get(
                              "mlp_dim", self.mlp_dim
                          ),
                          attention_fn=self.omninet.attention_fn,
                          attention_configs=self.omninet.attention_configs,
                          dropout_rate=self.dropout_rate,
                          attention_dropout_rate=self.attention_dropout_rate,
                          droplayer_p=droplayer_p,
                          post_sa_fn=post_sa_fn,
                          name=f"omni_encoder_block_{lyr}",
                      )(
                          xp, inputs_kv=None, deterministic=not train
                      )
                  elif self.omninet.encoder.type == "factorized":
                      if self.omninet.integrate != "factorized":
                          raise ValueError(
                              "omninet.encoder.type is 'factorized', "
                              f"but omninet.integrate is {self.omninet.integrate}."
                          )
                      xp = fast_vit_model_utils.EncoderAxialBlock(
                          mlp_dim=self.omninet.get(
                              "mlp_dim", self.mlp_dim
                          ),
                          attention_configs=self.omninet.attention_configs,
                          attention_fn=self.omninet.attention_fn,
                          factorization_axis=self.omninet.encoder.get(
                              "factorization_axis", (1,)
                          ),  # Apply attention only in depth.
                          dropout_rate=self.dropout_rate,
                          attention_dropout_rate=self.attention_dropout_rate,
                          droplayer_p=droplayer_p,
                          post_sa_fn=post_sa_fn,
                          name=f"omni_encoder_block_{lyr}",
                      )(xp, deterministic=not train)
                  else:
                      raise ValueError(
                          f"Unknown omninet.encoder.type: {self.omninet.encoder.type}"
                      )
  
                  if not self.omninet.pool_after_sa:
                      xp = _pool(xp)
              #################### TopQuery ########################
              elif self.omninet.query_type == "top":
                  if self.omninet.integrate == "concat":
                      xp = jnp.concatenate(partition_layers, 1)
                  else:
                      raise ValueError
  
                  xp = fast_vit_model_utils.Encoder1DBlock(
                      mlp_dim=self.omninet.get("mlp_dim", self.mlp_dim),
                      attention_fn=self.omninet.attention_fn,
                      attention_configs=self.omninet.attention_configs,
                      dropout_rate=self.dropout_rate,
                      attention_dropout_rate=self.attention_dropout_rate,
                      droplayer_p=droplayer_p,
                      name=f"omni_encoder_block_{lyr}",
                  )(x, xp, deterministic=not train)
              ########################################################
              x = x + xp
              partition_layers = []
              skip = True
  
      encoded = nn.LayerNorm(name="encoder_norm")(x)
      return encoded
  
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 OmnidirectionalEncoder1D obj's underlying member method code with example usage: obj.__call__(inputs, train) and returns: encoded, xp, nn, jnp