code: def test_all_binary_builtins():
      def _test(
          g,
          lhs,
          rhs,
          binary_op,
          reducer,
          partial,
          nid,
          broadcast="none",
      ):
          # initialize node/edge features with uniform(-1, 1)
          hu, hv, he = generate_feature(g, broadcast, binary_op)
          if binary_op == "div":
              # op = div
              # lhs range: [-1, 1]
              # rhs range: [1, 2]
              # result range: [-1, 1]
              if rhs == "u":
                  hu = (hu + 3) / 2
              elif rhs == "v":
                  hv = (hv + 3) / 2
              elif rhs == "e":
                  he = (he + 3) / 2
  
          if binary_op == "add" or binary_op == "sub":
              # op = add, sub
              # lhs range: [-1/2, 1/2]
              # rhs range: [-1/2, 1/2]
              # result range: [-1, 1]
              hu = hu / 2
              hv = hv / 2
              he = he / 2
  
          g.ndata["u"] = F.attach_grad(F.clone(hu))
          g.ndata["v"] = F.attach_grad(F.clone(hv))
          g.edata["e"] = F.attach_grad(F.clone(he))
  
          builtin_msg_name = "{}_{}_{}".format(lhs, binary_op, rhs)
          builtin_msg = getattr(fn, builtin_msg_name)
          builtin_red = getattr(fn, reducer)
  
          def target_feature_switch(g, target):
              if target == "u":
                  return g.ndata["u"]
              elif target == "v":
                  return g.ndata["v"]
              else:
                  return g.edata["e"]
  
          with F.record_grad():
              if partial:
                  g.pull(
                      nid,
                      builtin_msg(lhs, rhs, "m"),
                      builtin_red("m", "r1"),
                  )
              else:
                  g.update_all(
                      builtin_msg(lhs, rhs, "m"), builtin_red("m", "r1")
                  )
              r1 = g.ndata.pop("r1")
              F.backward(F.reduce_sum(r1))
              lhs_grad_1 = F.grad(target_feature_switch(g, lhs))
              rhs_grad_1 = F.grad(target_feature_switch(g, rhs))
  
          # reset grad
          g.ndata["u"] = F.attach_grad(F.clone(hu))
          g.ndata["v"] = F.attach_grad(F.clone(hv))
          g.edata["e"] = F.attach_grad(F.clone(he))
  
          def target_switch(edges, target):
              if target == "u":
                  return edges.src
              elif target == "v":
                  return edges.dst
              elif target == "e":
                  return edges.data
              else:
                  assert 0, "Unknown target {}".format(target)
  
          def mfunc(edges):
              op = getattr(F, binary_op)
              lhs_data = target_switch(edges, lhs)[lhs]
              rhs_data = target_switch(edges, rhs)[rhs]
              # NOTE(zihao): we need to do batched broadcast
              # e.g. (68, 3, 1) op (68, 5, 3, 4)
              while F.ndim(lhs_data) < F.ndim(rhs_data):
                  lhs_data = F.unsqueeze(lhs_data, 1)
              while F.ndim(rhs_data) < F.ndim(lhs_data):
                  rhs_data = F.unsqueeze(rhs_data, 1)
              return {"m": op(lhs_data, rhs_data)}
  
          def rfunc(nodes):
              op = getattr(F, reducer)
              return {"r2": op(nodes.mailbox["m"], 1)}
  
          with F.record_grad():
              if partial:
                  g.pull(nid, mfunc, rfunc)
              else:
                  g.update_all(mfunc, rfunc)
              r2 = g.ndata.pop("r2")
              F.backward(F.reduce_sum(r2), F.tensor([1.0]))
              lhs_grad_2 = F.grad(target_feature_switch(g, lhs))
              rhs_grad_2 = F.grad(target_feature_switch(g, rhs))
  
          rtol = 1e-4
          atol = 1e-4
  
          def _print_error(a, b):
              print(
                  "ERROR: Test {}_{}_{}_{} broadcast: {} partial: {}".format(
                      lhs, binary_op, rhs, reducer, broadcast, partial
                  )
              )
              return
              if lhs == "u":
                  lhs_data = hu
              elif lhs == "v":
                  lhs_data = hv
              elif lhs == "e":
                  lhs_data = he
  
              if rhs == "u":
                  rhs_data = hu
              elif rhs == "v":
                  rhs_data = hv
              elif rhs == "e":
                  rhs_data = he
              print("lhs", F.asnumpy(lhs_data).tolist())
              print("rhs", F.asnumpy(rhs_data).tolist())
              for i, (x, y) in enumerate(
                  zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())
              ):
                  if not np.allclose(x, y, rtol, atol):
                      print("@{} {} v.s. {}".format(i, x, y))
  
          if not F.allclose(r1, r2, rtol, atol):
              _print_error(r1, r2)
          assert F.allclose(r1, r2, rtol, atol)
  
          if not F.allclose(lhs_grad_1, lhs_grad_2, rtol, atol):
              print("left grad")
              _print_error(lhs_grad_1, lhs_grad_2)
          assert F.allclose(lhs_grad_1, lhs_grad_2, rtol, atol)
  
          if not F.allclose(rhs_grad_1, rhs_grad_2, rtol, atol):
              print("right grad")
              _print_error(rhs_grad_1, rhs_grad_2)
          assert F.allclose(rhs_grad_1, rhs_grad_2, rtol, atol)
  
      g = dgl.graph([])
      g.add_nodes(20)
      # NOTE(zihao): add self-loop to avoid zero-degree nodes.
      g.add_edges(g.nodes(), g.nodes())
      for i in range(2, 18):
          g.add_edges(0, i)
          g.add_edges(1, i)
          g.add_edges(i, 18)
          g.add_edges(i, 19)
      g.add_edges(18, 0)
      g.add_edges(18, 1)
      g.add_edges(19, 0)
      g.add_edges(19, 1)
      g = g.to(F.ctx())
      nid = F.tensor([0, 1, 4, 5, 7, 12, 14, 15, 18, 19], g.idtype)
      target = ["u", "v", "e"]
  
      for lhs, rhs in product(target, target):
          if lhs == rhs:
              continue
          for binary_op in ["add", "sub", "mul", "div"]:
              for reducer in ["sum", "max", "min", "mean"]:
                  for broadcast in ["none", lhs, rhs]:
                      for partial in [False, True]:
                          print(
                              lhs,
                              rhs,
                              binary_op,
                              reducer,
                              broadcast,
                              partial,
                          )
                          _test(
                              g,
                              lhs,
                              rhs,
                              binary_op,
                              reducer,
                              partial,
                              nid,
                              broadcast=broadcast,
                          )
  
file_path: dgl/tests/python/common/test_heterograph-kernel.py
imports: from itertools import product
  import backend
  import dgl
  import dgl.function
  import networkx
  import numpy
  import pytest
  from utils import get_cases, parametrize_idtype
label_desc: Use image node_id 9 for calling a global function with example usage: test_all_binary_builtins() and returns: dict, dict, g, edges, g, g, edges, edges