class_docstr: ## Squash
  
  This is **squashing** function from paper, given by equation $(1)$.
  
  $$\mathbf{v}_j = rac{{\lVert \mathbf{s}_j 
  Vert}^2}{1 + {\lVert \mathbf{s}_j 
  Vert}^2}
   rac{\mathbf{s}_j}{\lVert \mathbf{s}_j 
  Vert}$$
  
  $rac{\mathbf{s}_j}{\lVert \mathbf{s}_j 
  Vert}$
  normalizes the length of all the capsules, whilst
  $rac{{\lVert \mathbf{s}_j 
  Vert}^2}{1 + {\lVert \mathbf{s}_j 
  Vert}^2}$
  shrinks the capsules that have a length smaller than one .
class_docstr_tok: ["#","#","Squash","This","is","*","*","squashing","*","*","function","from","paper",",","given","by","equation","$","(","1",")","$",".","$","$","\\mathbf","{","v","}","_j","=","rac","{","{","\\lVert","\\mathbf","{","s","}","_j","Vert","}","^2","}","{","1","+","{","\\lVert","\\mathbf","{","s","}","_j","Vert","}","^2","}","rac","{","\\mathbf","{","s","}","_j","}","{","\\lVert","\\mathbf","{","s","}","_j","Vert","}","$","$","$","rac","{","\\mathbf","{","s","}","_j","}","{","\\lVert","\\mathbf","{","s","}","_j","Vert","}","$","normalizes","the","length","of","all","the","capsules",",","whilst","$","rac","{","{","\\lVert","\\mathbf","{","s","}","_j","Vert","}","^2","}","{","1","+","{","\\lVert","\\mathbf","{","s","}","_j","Vert","}","^2","}","$","shrinks","the","capsules","that","have","a","length","smaller","than","one","."]
docstr: The shape of `s` is `[batch_size, n_capsules, n_features]`
docstr_tok: ["The","shape","of","`","s","`","is","`","[","batch_size",",","n_capsules",",","n_features","]","`"]
code: def forward(self, s: torch.Tensor):
      """
      The shape of `s` is `[batch_size, n_capsules, n_features]`
      """
  
      # ${\lVert \mathbf{s}_j \rVert}^2$
      s2 = (s**2).sum(dim=-1, keepdims=True)
  
      # We add an epsilon when calculating $\lVert \mathbf{s}_j \rVert$ to make sure it doesn't become zero.
      # If this becomes zero it starts giving out `nan` values and training fails.
      # $$\mathbf{v}_j = \frac{{\lVert \mathbf{s}_j \rVert}^2}{1 + {\lVert \mathbf{s}_j \rVert}^2}
      # \frac{\mathbf{s}_j}{\sqrt{{\lVert \mathbf{s}_j \rVert}^2 + \epsilon}}$$
      return (s2 / (1 + s2)) * (s / torch.sqrt(s2 + self.epsilon))
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/capsule_networks/__init__.py
imports: import torch.nn
  import torch.nn.functional
  import torch.utils.data
  from labml_helpers.module import Module
label_desc: Use image node_id 2 for calling the Squash obj's underlying member method code with example usage: obj.forward(s) and returns: unknown