code: def test_csrmm_backward(idtype, dtype, num_vtypes):
      a, A = _random_simple_graph(
          idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB"
      )
      b, B = _random_simple_graph(
          idtype,
          dtype,
          F.ctx(),
          4,
          3,
          6,
          "B",
          "A" if num_vtypes == 1 else "C",
          "BA",
      )
      A_row, A_col = A.edges(order="eid")
      B_row, B_col = B.edges(order="eid")
      A_row = F.asnumpy(A_row)
      A_col = F.asnumpy(A_col)
      B_row = F.asnumpy(B_row)
      B_col = F.asnumpy(B_col)
      a_dense = F.attach_grad(F.tensor(a.todense(), dtype=dtype))
      b_dense = F.attach_grad(F.tensor(b.todense(), dtype=dtype))
  
      A.edata["w"] = F.attach_grad(A.edata["w"])
      B.edata["w"] = F.attach_grad(B.edata["w"])
  
      with F.record_grad():
          C = dgl.adj_product_graph(A, B, "w")
          assert len(C.ntypes) == num_vtypes
          assert len(C.etypes) == 1
          C_dense = np.zeros((3, 3))
          C_row, C_col = C.edges(order="eid")
          C_row = F.asnumpy(C_row)
          C_col = F.asnumpy(C_col)
          C_dense[C_row, C_col] = F.asnumpy(C.edata["w"])
          c_dense = F.matmul(a_dense, b_dense)
          assert np.allclose(
              C_dense, F.asnumpy(c_dense), rtol=1e-4, atol=1e-4
          )
  
          F.backward(F.reduce_sum(C.edata["w"]) + F.reduce_sum(c_dense))
          a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col]
          b_dense_grad = F.asnumpy(F.grad(b_dense))[B_row, B_col]
          A_spspmm_grad = F.asnumpy(F.grad(A.edata["w"]))
          B_spspmm_grad = F.asnumpy(F.grad(B.edata["w"]))
          assert np.allclose(
              a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4
          )
          assert np.allclose(
              b_dense_grad, B_spspmm_grad, rtol=1e-4, atol=1e-4
          )
  
file_path: dgl/tests/python/common/test_sparse_ops-csr.py
imports: import backend
  import dgl
  import numpy
  import pytest
  import scipy.sparse
  from utils import parametrize_idtype
label_desc: Use image node_id 3 for calling a global function with example usage: test_csrmm_backward(idtype, dtype, num_vtypes) without return types