code: def reference_implementation(
      hidden_states: torch.Tensor,
      weight: torch.Tensor,
      bias: Optional[torch.Tensor],
      act_type: ActivationType,
  ) -> torch.Tensor:
      dtype = hidden_states.dtype
      out_states = torch.nn.functional.linear(
          hidden_states, weight, bias
      )
      out_states.float()
  
      if is_gated(act_type):
          act_func_map = {
              ActivationType.ReGLU: torch.nn.functional.relu,
              ActivationType.GEGLU: lambda x: torch.nn.functional.gelu(
                  x, approximate="tanh"
              ),
              ActivationType.SiGLU: torch.nn.functional.silu,
          }
  
          act_act = out_states[..., ::2]
          act_linear = out_states[..., 1::2]
  
          act_act = act_func_map[act_type](act_act)
          out_states = act_act * act_linear
      else:
          act_func_map = {
              ActivationType.RELU: torch.nn.functional.relu,
              ActivationType.GELU: torch.nn.functional.gelu,
              ActivationType.SILU: torch.nn.functional.silu,
              ActivationType.IDENTITY: lambda x: x,
          }
  
          out_states = act_func_map[act_type](out_states)
      return out_states.to(dtype)
  
file_path: DeepSpeed/tests/unit/inference/v2/modules/test_blas_linear_module.py
imports: from typing import Optional
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum, is_gated
  from deepspeed.inference.v2.modules import ConfigBundle
  from deepspeed.inference.v2.modules.configs import DSLinearConfig
  from deepspeed.inference.v2.modules.interfaces import DSLinearRegistry
  from ...v2.inference_test_utils import allclose
label_desc: Use image node_id 1 for calling a global function with example usage: reference_implementation(hidden_states, weight, bias, act_type) and returns: out_states