code: def _blas_linear_helper(
      tokens: int,
      in_channels: int,
      out_channels: int,
      dtype: DtypeEnum,
      act_fn: ActivationType,
      use_bias: bool = True,
  ) -> None:
      linear_config = DSLinearConfig(
          max_tokens=2048,
          in_channels=in_channels,
          out_channels=out_channels,
          activation=act_fn,
          input_dtype=dtype,
          output_dtype=dtype,
      )
  
      bundle = ConfigBundle(name="blas_fp_linear", config=linear_config)
  
      module = DSLinearRegistry.instantiate_config(bundle)
  
      # Input vals
      hidden_states = (
          torch.randn(
              (tokens, in_channels),
              dtype=dtype.value,
              device=get_accelerator().current_device_name(),
          )
          * 0.01
      )
  
      weight_out_channels = (
          2 * out_channels if is_gated(act_fn) else out_channels
      )
      weight = (
          torch.randn(
              (weight_out_channels, in_channels),
              dtype=dtype.value,
              device=get_accelerator().current_device_name(),
          )
          * 0.01
      )
      if use_bias:
          bias = (
              torch.randn(
                  (weight_out_channels),
                  dtype=dtype.value,
                  device=get_accelerator().current_device_name(),
              )
              * 0.01
          )
      else:
          bias = None
  
      # Reference output
      ref_output = reference_implementation(
          hidden_states, weight, bias, act_fn
      )
  
      # New output
      ds_output = module(hidden_states, weight, bias)
  
      # Check
      assert allclose(ds_output, ref_output)
  
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 2 for calling a global function with example usage: _blas_linear_helper(tokens, in_channels, out_channels, dtype, act_fn, use_bias) without return types