code: def test_rms_norm(batch, seq_len, channels, dtype):
      device = get_accelerator().current_device_name()
      vals = torch.randn(
          (batch, seq_len, channels), dtype=dtype, device=device
      )
      gamma = torch.randn((channels), dtype=dtype, device=device)
      epsilon = 1e-5
  
      ref_output = ref_implementation(vals, gamma, epsilon)
      new_output = ds_implementation(vals, gamma, epsilon)
  
      assert allclose(new_output, ref_output)
  
file_path: deepspeed/tests/unit/ops/transformer/inference/test_rms_norm.py
imports: import deepspeed
  import torch
  import pytest
  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_rms_norm(batch, seq_len, channels, dtype) without return types