code: def test_custom_registration():
      channels = 4096
      dtype = torch.float16
      tokens = 1024
  
      config = DSNormConfig(
          max_tokens=2048,
          type="layer_norm",
          channels=channels,
          residual_dtype=dtype,
          input_dtype=dtype,
          output_dtype=dtype,
          eps=1e-5,
      )
      bundle = ConfigBundle(name="custom_post_ln", config=config)
  
      # Input vals
      hidden_states = torch.randn(
          (tokens, channels),
          dtype=dtype,
          device=get_accelerator().current_device_name(),
      )
      residual = torch.randn(
          (tokens, channels),
          dtype=dtype,
          device=get_accelerator().current_device_name(),
      )
      gamma = torch.randn(
          (channels),
          dtype=torch.float32,
          device=get_accelerator().current_device_name(),
      )
      beta = torch.rand(
          (channels),
          dtype=torch.float32,
          device=get_accelerator().current_device_name(),
      )
      epsilon = 1e-5
  
      # Reference output
      ref_output = reference_implementation(
          residual, hidden_states, gamma, beta, epsilon
      )
  
      # New output
      post_ln_module = DSPostNormRegistry.instantiate_config(bundle)
      gamma = post_ln_module.transform_param(gamma)
      beta = post_ln_module.transform_param(beta)
      ds_output, _ = post_ln_module(
          residual, hidden_states, gamma, beta
      )
  
      # Check
      assert allclose(ds_output, ref_output)
  
file_path: DeepSpeed/tests/unit/inference/v2/modules/test_custom_module.py
imports: import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.modules import ConfigBundle
  from deepspeed.inference.v2.modules.interfaces import DSPostNormRegistry
  from deepspeed.inference.v2.modules.configs import DSNormConfig
  from deepspeed.inference.v2.modules.implementations import cuda_post_ln
  from ...v2.inference_test_utils import allclose
label_desc: Use image node_id 2 for calling a global function with example usage: test_custom_registration() without return types