code: def test_spmv_3d_feat(idtype):
      def src_mul_edge_udf(edges):
          return {
              "sum": edges.src["h"]
              * F.unsqueeze(F.unsqueeze(edges.data["h"], 1), 1)
          }
  
      def sum_udf(nodes):
          return {"h": F.sum(nodes.mailbox["sum"], 1)}
  
      n = 100
      p = 0.1
      a = sp.random(n, n, p, data_rvs=lambda n: np.ones(n))
      g = dgl.from_scipy(a)
      g = g.astype(idtype).to(F.ctx())
      m = g.num_edges()
  
      # test#1: v2v with adj data
      h = F.randn((n, 5, 5))
      e = F.randn((m,))
  
      g.ndata["h"] = h
      g.edata["h"] = e
      g.update_all(
          message_func=fn.u_mul_e("h", "h", "sum"),
          reduce_func=fn.sum("sum", "h"),
      )  # 1
      ans = g.ndata["h"]
  
      g.ndata["h"] = h
      g.edata["h"] = e
      g.update_all(
          message_func=src_mul_edge_udf, reduce_func=fn.sum("sum", "h")
      )  # 2
      assert F.allclose(g.ndata["h"], ans)
  
      g.ndata["h"] = h
      g.edata["h"] = e
      g.update_all(
          message_func=src_mul_edge_udf, reduce_func=sum_udf
      )  # 3
      assert F.allclose(g.ndata["h"], ans)
  
      # test#2: e2v
      def src_mul_edge_udf(edges):
          return {"sum": edges.src["h"] * edges.data["h"]}
  
      h = F.randn((n, 5, 5))
      e = F.randn((m, 5, 5))
  
      g.ndata["h"] = h
      g.edata["h"] = e
      g.update_all(
          message_func=fn.u_mul_e("h", "h", "sum"),
          reduce_func=fn.sum("sum", "h"),
      )  # 1
      ans = g.ndata["h"]
  
      g.ndata["h"] = h
      g.edata["h"] = e
      g.update_all(
          message_func=src_mul_edge_udf, reduce_func=fn.sum("sum", "h")
      )  # 2
      assert F.allclose(g.ndata["h"], ans)
  
      g.ndata["h"] = h
      g.edata["h"] = e
      g.update_all(
          message_func=src_mul_edge_udf, reduce_func=sum_udf
      )  # 3
      assert F.allclose(g.ndata["h"], ans)
  
file_path: dgl/tests/python/common/test_heterograph-specialization.py
imports: import backend
  import dgl
  import dgl.function
  import numpy
  import scipy.sparse
  from utils import parametrize_idtype
label_desc: Use image node_id 7 for calling a global function with example usage: test_spmv_3d_feat(idtype) and returns: dict, dict, dict