class_docstr: The MLP-Mixer module.
  
  Args:
      num_tokens (int): Number of tokens/patches in each sample.
      in_channels (int): Input channels.
      out_channels (int): Output channels.
      dropout (float, optional): Dropout probability. (default: :obj:`0.0`)
docstr: Forward pass.
  
  Args:
      x (torch.Tensor): Tensor of size
          :obj:`[*, num_tokens, in_channels]`.
  
  Returns:
      Tensor of size :obj:`[*, out_channels]`.
docstr_tok: ["Forward","pass",".","Args",":","x","(","torch.Tensor",")",":","Tensor","of","size",":","obj",":","`","[","*",",","num_tokens",",","in_channels","]","`",".","Returns",":","Tensor","of","size",":","obj",":","`","[","*",",","out_channels","]","`","."]
code: def forward(self, x: Tensor) -> Tensor:
      r"""Forward pass.
  
      Args:
          x (torch.Tensor): Tensor of size
              :obj:`[*, num_tokens, in_channels]`.
  
      Returns:
          Tensor of size :obj:`[*, out_channels]`.
      """
      # Token mixing:
      h = self.token_norm(x).mT
      h = self.token_lin1(h)
      h = F.gelu(h)
      h = F.dropout(h, p=self.dropout, training=self.training)
      h = self.token_lin2(h)
      h = F.dropout(h, p=self.dropout, training=self.training)
      h_token = h.mT + x
  
      # Channel mixing:
      h = self.channel_norm(h_token)
      h = self.channel_lin1(h)
      h = F.gelu(h)
      h = F.dropout(h, p=self.dropout, training=self.training)
      h = self.channel_lin2(h)
      h = F.dropout(h, p=self.dropout, training=self.training)
      h_channel = h + h_token
  
      # Head:
      out = self.head_norm(h_channel)
      out = out.mean(dim=1)
      out = self.head_lin(out)
      return out
  
imports: import numpy
  import torch
  import torch.nn.functional
  from torch import Tensor
  from torch.nn import LayerNorm, Linear
  from torch_geometric.nn import TemporalEncoding
  from torch_geometric.utils import scatter, to_dense_batch
label_desc: Use image node_id 3 for calling the _MLPMixer obj's underlying member method code with example usage: obj.forward(x) and returns: out