code: def test_dtypes(shape: Iterable[int], dtype: torch.dtype) -> None:
      input_tensor = torch.randn(
          shape,
          dtype=dtype,
          device=get_accelerator().current_device_name(),
      )
  
      # Reference output
      ref_output = reference_geglu_implementation(
          input_tensor, act_fn=ActivationType.GEGLU
      )
  
      # Build kernel
      geglu = CUDAGatedActivation(
          input_tensor.size(-1),
          input_tensor.dtype,
          ActivationType.GEGLU,
      )
  
      # New output
      output_shape = list(input_tensor.shape)
      output_shape[-1] //= 2
      output_tensor = torch.empty(
          output_shape,
          dtype=input_tensor.dtype,
          device=get_accelerator().current_device_name(),
      )
      geglu(output_tensor, input_tensor)
  
      # Check
      assert allclose(output_tensor, ref_output)
  
file_path: deepspeed/tests/unit/inference/v2/kernels/core_ops/test_gated_activation.py
imports: from typing import Iterable, Optional
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.kernels.core_ops import CUDAGatedActivation
  from deepspeed.inference.v2.inference_utils import ActivationType
  from ....v2.inference_test_utils import get_dtypes, allclose
label_desc: Use image node_id 2 for calling a global function with example usage: test_dtypes(shape, dtype) without return types