code: def test_gather_mm_idx_b(feat_size, dtype, tol):
      if F._default_context_str == "cpu" and dtype == torch.float16:
          pytest.skip("float16 is not supported on CPU.")
      if (
          F._default_context_str == "gpu"
          and dtype == torch.bfloat16
          and not torch.cuda.is_bf16_supported()
      ):
          pytest.skip("BF16 is not supported.")
  
      if (
          F._default_context_str == "gpu"
          and dtype == torch.float16
          and torch.cuda.get_device_capability() < (7, 0)
      ):
          pytest.skip(
              f"FP16 is not supported for atomic operations on GPU with "
              f"cuda capability ({torch.cuda.get_device_capability()})."
          )
  
      dev = F.ctx()
      # input
      a = torch.tensor(np.random.rand(100, feat_size)).to(dev).to(dtype)
      a.requires_grad_()
      b = (
          torch.tensor(np.random.rand(10, feat_size, feat_size + 1))
          .to(dev)
          .to(dtype)
      )
      b.requires_grad_()
      idx = torch.tensor(np.random.randint(0, 10, 100)).to(dev).long()
      dc = (
          torch.tensor(np.random.rand(100, feat_size + 1))
          .to(dev)
          .to(dtype)
      )
      # compute
      c = gather_mm(a, b, idx_b=idx)
      c.backward(dc)
      da = a.grad.clone()
      db = b.grad.clone()
      # ground truth
      c_t = torch.bmm(a.unsqueeze(1), b[idx]).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=tol, rtol=tol)
      assert torch.allclose(da, da_t, atol=tol, rtol=tol)
      assert torch.allclose(db, db_t, atol=tol, rtol=tol)
  
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 9 for calling a global function with example usage: test_gather_mm_idx_b(feat_size, dtype, tol) without return types