code: def _test_gather_mm_idx_a(idtype, feat_size):
      # TODO(minjie): currently disabled due to bugs in the CUDA kernel. Need to fix it later.
      import torch
  
      dev = F.ctx()
      # input
      a = torch.tensor(np.random.rand(10, feat_size)).to(dev)
      a.requires_grad_()
      b = torch.tensor(
          np.random.rand(100, feat_size, feat_size + 1)
      ).to(dev)
      b.requires_grad_()
      idx = torch.tensor(np.random.randint(0, 10, 100)).to(dev)
      dc = torch.tensor(np.random.rand(100, feat_size + 1)).to(dev)
      # compute
      c = gather_mm(a, b, idx_a=idx)
      c.backward(dc)
      da = a.grad.clone()
      db = b.grad.clone()
      # ground truth
      c_t = torch.bmm(a[idx].unsqueeze(1), b).squeeze(1)
      a.grad.zero_()
      b.grad.zero_()
      c_t.backward(dc)
      da_t = a.grad
      db_t = b.grad
  
      assert torch.allclose(c, c_t, atol=1e-4, rtol=1e-4)
      assert torch.allclose(da, da_t, atol=1e-4, rtol=1e-4)
      assert torch.allclose(db, db_t, atol=1e-4, rtol=1e-4)
  
imports: import random
  import unittest
  import backend
  import dgl
  import numpy
  import pytest
  import torch
  from dgl.ops import gather_mm, gsddmm, gspmm, segment_reduce
  from utils import parametrize_idtype
  from utils.graph_cases import get_cases
label_desc: Use image node_id 10 for calling a global function with example usage: _test_gather_mm_idx_a(idtype, feat_size) without return types