class_docstr: ## MLP Mixer
  
  This module is a drop-in replacement for [self-attention layer](../mha.html).
  It transposes the input tensor before feeding it to the MLP and transposes back,
  so that the MLP is applied across the sequence dimension (across tokens or image patches) instead
  of the feature dimension.
class_docstr_tok: ["#","#","MLP","Mixer","This","module","is","a","drop-in","replacement","for","[","self-attention","layer","]","(","..","\/mha.html",")",".","It","transposes","the","input","tensor","before","feeding","it","to","the","MLP","and","transposes","back",",","so","that","the","MLP","is","applied","across","the","sequence","dimension","(","across","tokens","or","image","patches",")","instead","of","the","feature","dimension","."]
docstr: The [normal attention module](../mha.html) can be fed with different token embeddings for
  $       ext{query}$,$   ext{key}$, and $        ext{value}$ and a mask.
  
  We follow the same function signature so that we can replace it directly.
  
  For MLP mixing, $$x =   ext{query} =    ext{key} =      ext{value}$$ and masking is not possible.
  Shape of `query` (and `key` and `value`) is `[seq_len, batch_size, d_model]`.
docstr_tok: ["The","[","normal","attention","module","]","(","..","\/mha.html",")","can","be","fed","with","different","token","embeddings","for","$","ext","{","query","}","$",",","$","ext","{","key","}","$",",","and","$","ext","{","value","}","$","and","a","mask",".","We","follow","the","same","function","signature","so","that","we","can","replace","it","directly",".","For","MLP","mixing",",","$","$","x","=","ext","{","query","}","=","ext","{","key","}","=","ext","{","value","}","$","$","and","masking","is","not","possible",".","Shape","of","`","query","`","(","and","`","key","`","and","`","value","`",")","is","`","[","seq_len",",","batch_size",",","d_model","]","`","."]
code: def forward(
      self,
      query: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
      mask: Optional[torch.Tensor] = None,
  ):
      """
      The [normal attention module](../mha.html) can be fed with different token embeddings for
      $\text{query}$,$\text{key}$, and $\text{value}$ and a mask.
  
      We follow the same function signature so that we can replace it directly.
  
      For MLP mixing, $$x = \text{query} = \text{key} = \text{value}$$ and masking is not possible.
      Shape of `query` (and `key` and `value`) is `[seq_len, batch_size, d_model]`.
      """
  
      # $\text{query}$,$\text{key}$, and $\text{value}$ all should be the same
      assert query is key and key is value
      # MLP mixer doesn't support masking. i.e. all tokens will see all other token embeddings.
      assert mask is None
  
      # Assign to `x` for clarity
      x = query
  
      # Transpose so that the last dimension is the sequence dimension.
      # New shape is `[d_model, batch_size, seq_len]`
      x = x.transpose(0, 2)
      # Apply the MLP across tokens
      x = self.mlp(x)
      # Transpose back into original form
      x = x.transpose(0, 2)
  
      #
      return x
  
imports: from typing import Optional
  import torch
  from torch import nn
label_desc: Use image node_id 2 for calling the MLPMixer obj's underlying member method code with example usage: obj.forward(query, key, value, mask) and returns: x