code: def AX(graph, X):
      """Y = AX"""
  
      graph.srcdata["h"] = X
      graph.update_all(
          fn.u_mul_e("h", "w", "m"),
          fn.sum("m", "h"),
      )
      Y = graph.dstdata["h"]
  
      return Y
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  from ....None import function
label_desc: Use image node_id 2 for calling a global function with example usage: AX(graph, X) and returns: Y