class_docstr: ## Routing Algorithm
  
  This is the routing mechanism described in the paper.
  You can use multiple routing layers in your models.
  
  This combines calculating $\mathbf{s}_j$ for this layer and
  the routing algorithm described in *Procedure 1*.
class_docstr_tok: ["#","#","Routing","Algorithm","This","is","the","routing","mechanism","described","in","the","paper",".","You","can","use","multiple","routing","layers","in","your","models",".","This","combines","calculating","$","\\mathbf","{","s","}","_j","$","for","this","layer","and","the","routing","algorithm","described","in","*","Procedure","1","*","."]
docstr: The shape of `u` is `[batch_size, n_capsules, n_features]`.
  These are the capsules from the lower layer.
docstr_tok: ["The","shape","of","`","u","`","is","`","[","batch_size",",","n_capsules",",","n_features","]","`",".","These","are","the","capsules","from","the","lower","layer","."]
code: def forward(self, u: torch.Tensor):
      """
      The shape of `u` is `[batch_size, n_capsules, n_features]`.
      These are the capsules from the lower layer.
      """
  
      # $$\hat{\mathbf{u}}_{j|i} = \mathbf{W}_{ij} \mathbf{u}_i$$
      # Here $j$ is used to index capsules in this layer, whilst $i$ is
      # used to index capsules in the layer below (previous).
      u_hat = torch.einsum("ijnm,bin->bijm", self.weight, u)
  
      # Initial logits $b_{ij}$ are the log prior probabilities that capsule $i$
      # should be coupled with $j$.
      # We initialize these at zero
      b = u.new_zeros(u.shape[0], self.in_caps, self.out_caps)
  
      v = None
  
      # Iterate
      for i in range(self.iterations):
          # routing softmax $$c_{ij} = \frac{\exp({b_{ij}})}{\sum_k\exp({b_{ik}})}$$
          c = self.softmax(b)
          # $$\mathbf{s}_j = \sum_i{c_{ij} \hat{\mathbf{u}}_{j|i}}$$
          s = torch.einsum("bij,bijm->bjm", c, u_hat)
          # $$\mathbf{v}_j = squash(\mathbf{s}_j)$$
          v = self.squash(s)
          # $$a_{ij} = \mathbf{v}_j \cdot \hat{\mathbf{u}}_{j|i}$$
          a = torch.einsum("bjm,bijm->bij", v, u_hat)
          # $$b_{ij} \gets b_{ij} + \mathbf{v}_j \cdot \hat{\mathbf{u}}_{j|i}$$
          b = b + a
  
      return v
  
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 Router obj's underlying member method code with example usage: obj.forward(u) and returns: v