class_docstr: Mixer block consisting of a token- and a channel-mixing phase.
  
  Attributes:
    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: The layer dropout rate (= stochastic depth).
    layer_scale: The scalar value used to initialise layer_scale. If None,
      layer_scale is not used.
  
  Returns:
    Output after mixer block.
class_docstr_tok: ["Mixer","block","consisting","of","a","token-","and","a","channel-mixing","phase",".","Attributes",":","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",":","The","layer","dropout","rate","(","=","stochastic","depth",")",".","layer_scale",":","The","scalar","value","used","to","initialise","layer_scale",".","If","None",",","layer_scale","is","not","used",".","Returns",":","Output","after","mixer","block","."]
docstr: Applies the Mixer block to inputs.
docstr_tok: ["Applies","the","Mixer","block","to","inputs","."]
code: def __call__(
      self, inputs: jnp.ndarray, deterministic: bool
  ) -> jnp.ndarray:
      """Applies the Mixer block to inputs."""
      if inputs.ndim != 3:
          raise ValueError(
              "Input should be of shape `[batch, tokens, channels]`."
          )
  
      if self.layer_scale is not None:
          layerscale_init = nn_layers.get_constant_initializer(
              self.layer_scale
          )
  
      # Token mixing part, provides between-patches communication.
      x = nn.LayerNorm()(inputs)
      x = jnp.swapaxes(x, 1, 2)
  
      x = attention_layers.MlpBlock(
          mlp_dim=self.sequence_mlp_dim,
          dropout_rate=self.dropout_rate,
          activation_fn=nn.gelu,
          name="token_mixing",
      )(x, deterministic=deterministic)
      x = jnp.swapaxes(x, 1, 2)
      if self.layer_scale is not None:
          x = nn_layers.Affine(
              scale_init=layerscale_init, use_bias=False
          )(x)
  
      x = nn_layers.StochasticDepth(rate=self.stochastic_depth)(
          x, deterministic
      )
      x = self.combine_branches(x, inputs)
  
      # Channel-mixing part, which provides within-patch communication.
      y = nn.LayerNorm()(x)
      y = attention_layers.MlpBlock(
          mlp_dim=self.channels_mlp_dim,
          dropout_rate=self.dropout_rate,
          activation_fn=nn.gelu,
          name="channel_mixing",
      )(y, deterministic=deterministic)
      if self.layer_scale is not None:
          x = nn_layers.Affine(
              scale_init=layerscale_init, use_bias=False
          )(x)
  
      y = nn_layers.StochasticDepth(rate=self.stochastic_depth)(
          y, deterministic
      )
      return self.combine_branches(y, 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 2 for calling the MixerBlock obj's underlying member method code with example usage: obj.__call__(inputs, deterministic) and returns: self