class_docstr: ## gMLP Block
  
  Each block does the following transformations to input embeddings
  $X \in \mathbb{R}^{n        imes d}$ where $n$ is the sequence length
  and $d$ is the dimensionality of the embeddings:
  
  egin{align}
  Z &= \sigma(XU) \
      ilde{Z} &= s(Z) \
  Y &=        ilde{Z}V \
  \end{align}
  
  where $V$ and $U$ are learnable projection weights.
  $s(\cdot)$ is the Spacial Gating Unit defined below.
  Output dimensionality of $s(\cdot)$ will be half of $Z$.
  $\sigma$ is an activation function such as
  [GeLU](https://pytorch.org/docs/stable/generated/torch.nn.GELU.html).
class_docstr_tok: ["#","#","gMLP","Block","Each","block","does","the","following","transformations","to","input","embeddings","$","X","\\in","\\mathbb","{","R","}","^","{","n","imes","d","}","$","where","$","n","$","is","the","sequence","length","and","$","d","$","is","the","dimensionality","of","the","embeddings",":","\begin","{","align","}","Z","&","=","\\sigma","(","XU",")","\\","ilde","{","Z","}","&","=","s","(","Z",")","\\","Y","&","=","ilde","{","Z","}","V","\\","\\end","{","align","}","where","$","V","$","and","$","U","$","are","learnable","projection","weights",".","$","s","(","\\cdot",")","$","is","the","Spacial","Gating","Unit","defined","below",".","Output","dimensionality","of","$","s","(","\\cdot",")","$","will","be","half","of","$","Z","$",".","$","\\sigma","$","is","an","activation","function","such","as","[","GeLU","]","(","https",":","\/\/pytorch.org\/docs\/stable\/generated\/torch.nn.GELU.html",")","."]
docstr: * `d_model` is the dimensionality ($d$) of $X$
  * `d_ffn` is the dimensionality of $Z$
  * `seq_len` is the length of the token sequence ($n$)
docstr_tok: ["*","`","d_model","`","is","the","dimensionality","(","$","d","$",")","of","$","X","$","*","`","d_ffn","`","is","the","dimensionality","of","$","Z","$","*","`","seq_len","`","is","the","length","of","the","token","sequence","(","$","n","$",")"]
code: def __init__(self, d_model: int, d_ffn: int, seq_len: int):
      """
      * `d_model` is the dimensionality ($d$) of $X$
      * `d_ffn` is the dimensionality of $Z$
      * `seq_len` is the length of the token sequence ($n$)
      """
      super().__init__()
      # Normalization layer fro Pre-Norm
      self.norm = nn.LayerNorm([d_model])
      # Activation function $\sigma$
      self.activation = nn.GELU()
      # Projection layer for $Z = \sigma(XU)$
      self.proj1 = nn.Linear(d_model, d_ffn)
      # Spacial Gating Unit $s(\cdot)$
      self.sgu = SpacialGatingUnit(d_ffn, seq_len)
      # Projection layer for $Y = \tilde{Z}V$
      self.proj2 = nn.Linear(d_ffn // 2, d_model)
      # Embedding size (required by [Encoder](../models.html#Encoder).
      # We use the encoder module from transformer architecture and plug
      # *gMLP* block as a replacement for the [Transformer Layer](../models.html#Encoder).
      self.size = d_model
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/gmlp/__init__.py
imports: from typing import Optional
  import torch
  from torch import nn
label_desc: Use image node_id 1 to create a new GMLPBlock object from inherited base classes: nn with example: obj = GMLPBlock(d_model, d_ffn, seq_len)