code: def forward(self, u_hat, routing_num=1):
      self.g.edata["u_hat"] = u_hat
      batch_size = self.batch_size
  
      # step 2 (line 5)
      def cap_message(edges):
          if batch_size:
              return {
                  "m": edges.data["c"].unsqueeze(1)
                  * edges.data["u_hat"]
              }
          else:
              return {"m": edges.data["c"] * edges.data["u_hat"]}
  
      def cap_reduce(nodes):
          return {"s": th.sum(nodes.mailbox["m"], dim=1)}
  
      for r in range(routing_num):
          # step 1 (line 4): normalize over out edges
          edges_b = self.g.edata["b"].view(
              self.in_nodes, self.out_nodes
          )
          self.g.edata["c"] = F.softmax(edges_b, dim=1).view(-1, 1)
  
          # Execute step 1 & 2
          self.g.update_all(
              message_func=cap_message, reduce_func=cap_reduce
          )
  
          # step 3 (line 6)
          if self.batch_size:
              self.g.nodes[self.out_indx].data["v"] = squash(
                  self.g.nodes[self.out_indx].data["s"], dim=2
              )
          else:
              self.g.nodes[self.out_indx].data["v"] = squash(
                  self.g.nodes[self.out_indx].data["s"], dim=1
              )
          # step 4 (line 7)
          v = th.cat(
              [self.g.nodes[self.out_indx].data["v"]] * self.in_nodes,
              dim=0,
          )
          if self.batch_size:
              self.g.edata["b"] = self.g.edata["b"] + (
                  self.g.edata["u_hat"] * v
              ).mean(dim=1).sum(dim=1, keepdim=True)
          else:
              self.g.edata["b"] = self.g.edata["b"] + (
                  self.g.edata["u_hat"] * v
              ).sum(dim=1, keepdim=True)
  
file_path: dgl/examples/pytorch/capsule/DGLRoutingLayer.py
imports: import dgl
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the DGLRoutingLayer obj's underlying member method code with example usage: obj.forward(u_hat, routing_num) and returns: dict, dict, dict