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: * `x` is the input embedding tensor $X$ of shape `[seq_len, batch_size, d_model]`
  * `mask` is a boolean mask of shape `[seq_len, seq_len, 1]` that controls the visibility of tokens
   among each other.
docstr_tok: ["*","`","x","`","is","the","input","embedding","tensor","$","X","$","of","shape","`","[","seq_len",",","batch_size",",","d_model","]","`","*","`","mask","`","is","a","boolean","mask","of","shape","`","[","seq_len",",","seq_len",",","1","]","`","that","controls","the","visibility","of","tokens","among","each","other","."]
code: def forward(
      self, *, x: torch.Tensor, mask: Optional[torch.Tensor] = None
  ):
      """
      * `x` is the input embedding tensor $X$ of shape `[seq_len, batch_size, d_model]`
      * `mask` is a boolean mask of shape `[seq_len, seq_len, 1]` that controls the visibility of tokens
       among each other.
      """
      # Keep a copy for shortcut connection
      shortcut = x
      # Normalize $X$
      x = self.norm(x)
      # Projection and activation $Z = \sigma(XU)$
      z = self.activation(self.proj1(x))
      # Spacial Gating Unit $\tilde{Z} = s(Z)$
      z = self.sgu(z, mask)
      # Final projection $Y = \tilde{Z}V$
      z = self.proj2(z)
  
      # Add the shortcut connection
      return z + shortcut
  
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 2 for calling the GMLPBlock obj's underlying member method code with example usage: obj.forward() and returns: unknown