class_docstr: Mixer 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.
    dropout_rate: Dropout rate.
    stochastic_depth: overall stochastic depth rate.
class_docstr_tok: ["Mixer","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",".","dropout_rate",":","Dropout","rate",".","stochastic_depth",":","overall","stochastic","depth","rate","."]
code: def __call__(
      self, x: jnp.ndarray, *, train: bool, debug: bool = False
  ) -> jnp.ndarray:
      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])
      for i in range(self.num_layers):
          p = (i / max(self.num_layers - 1, 1)) * self.stochastic_depth
          x = MixerBlock(
              channels_mlp_dim=self.channels_mlp_dim,
              sequence_mlp_dim=self.sequence_mlp_dim,
              dropout_rate=self.dropout_rate,
              stochastic_depth=p,
              layer_scale=self.layer_scale,
              name=f"mixerblock_{i}",
          )(x, deterministic=not train)
      x = nn.LayerNorm(name="pre_logits_norm")(x)
      # Use global average pooling for classifier:
      x = jnp.mean(x, axis=1)
      x = nn_layers.IdentityLayer(name="pre_logits")(x)
      return nn.Dense(
          self.num_classes,
          kernel_init=nn.initializers.zeros,
          name="output_projection",
      )(x)
  
file_path: scenic/scenic/projects/baselines/mixer.py
imports: from typing import Optional, Sequence
  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
label_desc: Use image node_id 1 for calling the Mixer obj's underlying member method code with example usage: obj.__call__(x) and returns: nn