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: * `z` is the input $Z$ of shape `[seq_len, batch_size, d_z]`
  * `mask` is is a boolean mask of shape `[seq_len, seq_len, 1]` that controls the visibility of tokens
   among each other. The last dimension of size `1` is the batch, which we have in other transformer
   implementations and was left for compatibility.
docstr_tok: ["*","`","z","`","is","the","input","$","Z","$","of","shape","`","[","seq_len",",","batch_size",",","d_z","]","`","*","`","mask","`","is","is","a","boolean","mask","of","shape","`","[","seq_len",",","seq_len",",","1","]","`","that","controls","the","visibility","of","tokens","among","each","other",".","The","last","dimension","of","size","`","1","`","is","the","batch",",","which","we","have","in","other","transformer","implementations","and","was","left","for","compatibility","."]
code: def forward(
      self, z: torch.Tensor, mask: Optional[torch.Tensor] = None
  ):
      """
      * `z` is the input $Z$ of shape `[seq_len, batch_size, d_z]`
      * `mask` is is a boolean mask of shape `[seq_len, seq_len, 1]` that controls the visibility of tokens
       among each other. The last dimension of size `1` is the batch, which we have in other transformer
       implementations and was left for compatibility.
      """
  
      # Get sequence length
      seq_len = z.shape[0]
      # Split $Z$ into $Z_1$ and $Z_2$
      z1, z2 = torch.chunk(z, 2, dim=-1)
  
      # Check mask
      if mask is not None:
          # `mask` has shape `[seq_len_q, seq_len_k, batch_size]`.
          # The batch dimension should be of size `1` because this implementation supports
          # only same mask for all samples in the batch.
          assert mask.shape[0] == 1 or mask.shape[0] == seq_len
          assert mask.shape[1] == seq_len
          # Here we only support the same mask for all samples
          assert mask.shape[2] == 1
          # Remove the batch dimension
          mask = mask[:, :, 0]
  
      # Normalize $Z_2$ before $f_{W,b}(\cdot)$
      z2 = self.norm(z2)
      # Get the weight matrix; truncate if larger than `seq_len`
      weight = self.weight[:seq_len, :seq_len]
      # Apply mask to the weights.
      #
      # If $W_{i,j}$ is $0$ then $f_{W,b}(Z_2)_i$ will not get any information
      # from token $j$.
      if mask is not None:
          weight = weight * mask
  
      # $f_{W,b}(Z_2) = W Z_2 + b$
      z2 = (
          torch.einsum("ij,jbd->ibd", weight, z2)
          + self.bias[:seq_len, None, None]
      )
  
      # $Z_1 \odot f_{W,b}(Z_2)$
      return z1 * z2
  
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 SpacialGatingUnit obj's underlying member method code with example usage: obj.forward(z, mask) and returns: unknown