code: def test_moe_residual_matmul(hidden_dim, c, dtype):
      residual_ds = torch.randn(
          (c, hidden_dim * c, hidden_dim),
          dtype=dtype,
          device=get_accelerator().device_name(),
      )
      coeff1 = torch.randn(
          (1, 1, hidden_dim),
          dtype=dtype,
          device=get_accelerator().device_name(),
      )
      coeff2 = torch.randn(
          (1, 1, hidden_dim),
          dtype=dtype,
          device=get_accelerator().device_name(),
      )
      out_ds = torch.randn(
          (c, hidden_dim * c, hidden_dim),
          dtype=dtype,
          device=get_accelerator().device_name(),
      )
      coeff_ds = torch.cat((coeff1, coeff2), dim=-1)
      residual_ref = residual_ds.clone().detach()
      coeff_ref = coeff_ds.clone().detach()
      out_ref = out_ds.clone().detach()
  
      ds_out = run_moe_res_matmul_ds(residual_ds, coeff_ds, out_ds)
      ref_out = run_moe_res_matmul_reference(
          residual_ref, coeff1, coeff2, out_ref
      )
  
      assert allclose(ds_out, ref_out)
  
file_path: deepspeed/tests/unit/ops/transformer/inference/test_moe_res_matmult.py
imports: import pytest
  import torch
  import deepspeed
  from deepspeed.accelerator import get_accelerator
  from deepspeed.ops.op_builder import InferenceBuilder
  from .inference_test_utils import allclose, get_dtypes
label_desc: Use image node_id 3 for calling a global function with example usage: test_moe_residual_matmul(hidden_dim, c, dtype) without return types