code: def forward(self, u_hat, routing_num=1):
      self.g.edata["u_hat"] = u_hat
  
      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)
          self.g.edata["c u_hat"] = (
              self.g.edata["c"] * self.g.edata["u_hat"]
          )
  
          # Execute step 1 & 2
          self.g.update_all(fn.copy_e("c u_hat", "m"), fn.sum("m", "s"))
  
          # step 3 (line 6)
          self.g.nodes[self.out_indx].data["v"] = self.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,
          )
          self.g.edata["b"] = self.g.edata["b"] + (
              self.g.edata["u_hat"] * v
          ).sum(dim=1, keepdim=True)
  
file_path: dgl/tutorials/models/4_old_wines/2_capsule.py
imports: import os
  import dgl
  import matplotlib.pyplot
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import dgl.function
  import matplotlib.animation
  import seaborn
  import networkx
  from networkx.algorithms import bipartite
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) without return types