code: def test_residual_add(
      inference_module,
      batch,
      sequence,
      hidden_dim,
      dtype,
      mlp_after_attn,
      add_bias,
      mp_size,
      pre_attn_norm,
  ):
      ds_out = torch.randn(
          (batch, sequence, hidden_dim), dtype=dtype, device="cuda"
      )
      residual = torch.randn(
          (batch, sequence, hidden_dim), dtype=dtype, device="cuda"
      )
      attn_output = torch.randn(
          (batch, sequence, hidden_dim), dtype=dtype, device="cuda"
      )
      final_bias = torch.randn((hidden_dim), dtype=dtype, device="cuda")
      attn_bias = torch.randn((hidden_dim), dtype=dtype, device="cuda")
  
      ref_out = ds_out.clone()
      ref_out = run_residual_add_reference(
          ref_out,
          residual,
          attn_output,
          attn_bias,
          final_bias,
          mlp_after_attn,
          add_bias,
          mp_size,
          pre_attn_norm,
      )
  
      res_add_args = [
          ds_out,
          residual,
          attn_output,
          attn_bias,
          final_bias,
          mp_size,
          mlp_after_attn,
          add_bias,
          pre_attn_norm,
      ]
  
      if dtype == torch.float16:
          ds_out = inference_module.residual_add_bias_fp16(
              *res_add_args
          )
      elif dtype == torch.float32:
          ds_out = inference_module.residual_add_bias_fp32(
              *res_add_args
          )
      else:
          raise ValueError(f"Unsupported dtype: {dtype}")
  
      assert allclose(ds_out, ref_out)
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/tests/unit/ops/transformer/inference/test_residual_add.py
imports: import pytest
  import torch
  import deepspeed
  from deepspeed.ops.op_builder import InferenceBuilder
label_desc: Use image node_id 6 for calling a global function with example usage: test_residual_add(inference_module, batch, sequence, hidden_dim, dtype, mlp_after_attn, add_bias, mp_size, pre_attn_norm) without return types