code: def test_layer_norm(batch, seq_len, channels, dtype):
      vals = torch.randn(
          (batch, seq_len, channels),
          dtype=dtype,
          device=torch.cuda.current_device(),
      )
      gamma = torch.randn(
          (channels), dtype=dtype, device=torch.cuda.current_device()
      )
      beta = torch.rand(
          (channels), dtype=dtype, device=torch.cuda.current_device()
      )
      epsilon = 1e-5
  
      ref_output = ref_implementation(
          vals, gamma, beta, epsilon, channels, dtype
      )
      new_output = ds_implementation(vals, gamma, beta, epsilon)
  
      assert allclose(new_output, ref_output)
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/tests/unit/ops/transformer/inference/test_layer_norm.py
imports: import deepspeed
  import torch
  import pytest
  from deepspeed.ops.op_builder import InferenceBuilder
label_desc: Use image node_id 4 for calling a global function with example usage: test_layer_norm(batch, seq_len, channels, dtype) without return types