code: def compute_uhat(self, x):
      # x is the input vextor with shape [batch_size, in_nodes_dim, in_nodes]
      # Transpose x to [batch_size, in_nodes, in_nodes_dim]
      x = x.transpose(1, 2)
      # Expand x to [batch_size, in_nodes, out_nodes, in_nodes_dim, 1]
      x = torch.stack([x] * self.out_nodes, dim=2).unsqueeze(4)
      # Expand W from [in_nodes, out_nodes, in_nodes_dim, out_nodes_dim]
      # to [batch_size, in_nodes, out_nodes, out_nodes_dim, in_nodes_dim]
      W = self.weight.expand(self.batch_size, *self.weight.size())
      # u_hat's shape is [in_nodes, out_nodes, batch_size, out_nodes_dim]
      u_hat = (
          torch.matmul(W, x)
          .permute(1, 2, 0, 3, 4)
          .squeeze()
          .contiguous()
      )
      return u_hat.view(-1, self.batch_size, self.out_nodes_dim)
  
file_path: dgl/examples/pytorch/capsule/DGLDigitCapsule.py
imports: import dgl
  import dgl.function
  import torch
  from DGLRoutingLayer import DGLRoutingLayer
  from torch import nn
  from torch.nn import functional
label_desc: Use image node_id 3 for calling the DGLDigitCapsuleLayer obj's underlying member method code with example usage: obj.compute_uhat(x) and returns: u_hat