class_docstr: ## Spatial Gating Unit
  
  $$s(Z) = Z_1 \odot f_{W,b}(Z_2)$$
  
  where $f_{W,b}(Z) = W Z + b$ is a linear transformation along the sequence dimension,
  and $\odot$ is element-wise multiplication.
  $Z$ is split into to parts of equal size $Z_1$ and $Z_2$ along the channel dimension (embedding dimension).
class_docstr_tok: ["#","#","Spatial","Gating","Unit","$","$","s","(","Z",")","=","Z_1","\\odot","f_","{","W",",","b","}","(","Z_2",")","$","$","where","$","f_","{","W",",","b","}","(","Z",")","=","W","Z","+","b","$","is","a","linear","transformation","along","the","sequence","dimension",",","and","$","\\odot","$","is","element-wise","multiplication",".","$","Z","$","is","split","into","to","parts","of","equal","size","$","Z_1","$","and","$","Z_2","$","along","the","channel","dimension","(","embedding","dimension",")","."]
docstr: * `d_z` is the dimensionality of $Z$
  * `seq_len` is the sequence length
docstr_tok: ["*","`","d_z","`","is","the","dimensionality","of","$","Z","$","*","`","seq_len","`","is","the","sequence","length"]
code: def __init__(self, d_z: int, seq_len: int):
      """
      * `d_z` is the dimensionality of $Z$
      * `seq_len` is the sequence length
      """
      super().__init__()
      # Normalization layer before applying $f_{W,b}(\cdot)$
      self.norm = nn.LayerNorm([d_z // 2])
      # Weight $W$ in $f_{W,b}(\cdot)$.
      #
      # The paper notes that it's important to initialize weights to small values and the bias to $1$,
      # so that during the initial training $s(\cdot)$ is close to identity (apart from the split).
      self.weight = nn.Parameter(
          torch.zeros(seq_len, seq_len).uniform_(-0.01, 0.01),
          requires_grad=True,
      )
      # Weight $b$ in $f_{W,b}(\cdot)$
      #
      # The paper notes that it's important to initialize bias to $1$.
      self.bias = nn.Parameter(torch.ones(seq_len), requires_grad=True)
  
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 SpacialGatingUnit object from inherited base classes: nn with example: obj = SpacialGatingUnit(d_z, seq_len)