class_docstr: ## FNet - Mix tokens
  
  This module simply implements
  $$
  \mathcal{R}ig(\mathcal{F}_ ext{seq} ig(\mathcal{F}_       ext{hidden} (x) ig) ig)
  $$
  
  The structure of this module is made similar to a [standard attention module](../mha.html) so that we can simply
  replace it.
class_docstr_tok: ["#","#","FNet","-","Mix","tokens","This","module","simply","implements","$","$","\\mathcal","{","R","}","\big","(","\\mathcal","{","F","}","_","ext","{","seq","}","\big","(","\\mathcal","{","F","}","_","ext","{","hidden","}","(","x",")","\big",")","\big",")","$","$","The","structure","of","this","module","is","made","similar","to","a","[","standard","attention","module","]","(","..","\/mha.html",")","so","that","we","can","simply","replace","it","."]
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 FNet 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","FNet","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 FNet 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 equal to $x$ for token mixing
      assert query is key and key is value
      # Token mixing 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
  
      # Apply the Fourier transform along the hidden (embedding) dimension
      # $$\mathcal{F}_\text{hidden} (x)$$
      #
      # The output of the Fourier transform is a tensor of
      # [complex numbers](https://pytorch.org/docs/stable/complex_numbers.html).
      fft_hidden = torch.fft.fft(x, dim=2)
      # Apply the Fourier transform along the sequence dimension
      # $$\mathcal{F}_\text{seq} \big(\mathcal{F}_\text{hidden} (x) \big)$$
      fft_seq = torch.fft.fft(fft_hidden, dim=0)
  
      # Get the real component
      # $$\mathcal{R}\big(\mathcal{F}_\text{seq} \big(\mathcal{F}_\text{hidden} (x) \big) \big)$$
      return torch.real(fft_seq)
  
imports: from typing import Optional
  import torch
  from torch import nn
label_desc: Use image node_id 1 for calling the FNetMix obj's underlying member method code with example usage: obj.forward(query, key, value, mask) and returns: torch