code: def reference_implementation(
      residual: torch.Tensor,
      hidden_states: Optional[torch.Tensor],
      gamma: torch.Tensor,
      epsilon: float,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      dtype = residual.dtype
  
      if hidden_states is not None:
          hidden_states = hidden_states
          residual = residual + hidden_states
  
      rms_vals = residual.to(torch.float32)
      variance = rms_vals.pow(2).mean(-1, keepdim=True)
      rms_vals = rms_vals * torch.rsqrt(variance + epsilon)
  
      if gamma.dtype in [torch.float16, torch.bfloat16]:
          rms_vals = rms_vals.to(gamma.dtype)
  
      hidden_states = gamma * rms_vals
  
      return residual.to(dtype), hidden_states.to(dtype)
  
file_path: deepspeed/tests/unit/inference/v2/modules/test_pre_rms_module.py
imports: from typing import Optional, Tuple
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.modules import ConfigBundle
  from deepspeed.inference.v2.modules.configs import DSNormConfig
  from deepspeed.inference.v2.modules.interfaces import DSPreNormRegistry
  from ...v2.inference_test_utils import get_dtypes, allclose
label_desc: Use image node_id 1 for calling a global function with example usage: reference_implementation(residual, hidden_states, gamma, epsilon) and returns: residual,  hidden_states