code: def reference_implementation(
      residual: torch.Tensor,
      hidden_states: torch.Tensor,
      gamma: torch.Tensor,
      beta: torch.Tensor,
      epsilon: float,
  ) -> torch.Tensor:
      residual_f = residual.to(torch.float32)
      hidden_states_f = hidden_states.to(torch.float32)
      gamma_f = gamma.to(torch.float32)
      beta_f = beta.to(torch.float32)
      return torch.nn.functional.layer_norm(
          residual_f + hidden_states_f,
          (hidden_states_f.size(-1),),
          weight=gamma_f,
          bias=beta_f,
          eps=epsilon,
      ).to(hidden_states.dtype)
  
file_path: deepspeed/tests/unit/inference/v2/modules/test_custom_module.py
imports: import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.modules import ConfigBundle
  from deepspeed.inference.v2.modules.interfaces import DSPostNormRegistry
  from deepspeed.inference.v2.modules.configs import DSNormConfig
  from deepspeed.inference.v2.modules.implementations import cuda_post_ln
  from ...v2.inference_test_utils import allclose
label_desc: Use image node_id 1 for calling a global function with example usage: reference_implementation(residual, hidden_states, gamma, beta, epsilon) and returns: torch