code: def check_layernorm(lazy_init: bool):
      ctx = LazyInitContext() if lazy_init else nullcontext()
  
      norm = nn.LayerNorm(128, 0.00001).cuda()
      with ctx:
          norm_copy = nn.LayerNorm(128, 0.00001).cuda()
      norm1d = FusedLayerNorm.from_native_module(
          norm_copy, process_group=None
      )
  
      assert norm1d.weight.shape == torch.Size([128])
      assert norm_copy.weight is norm1d.weight
      assert norm_copy.bias is norm1d.bias
  
      # ensure state dict is reversibly loadable
      norm.load_state_dict(norm1d.state_dict())
      norm1d.load_state_dict(norm.state_dict())
  
      # check computation correctness
      x = torch.rand(4, 128).cuda()
      out = norm(x)
      gather_out = norm1d(x)
      assert_close(out, gather_out)
  
      # check backward correctness
      out.sum().backward()
      gather_out.sum().backward()
  
      assert_close(norm.weight.grad, norm1d.weight.grad)
  
file_path: ColossalAI/tests/test_shardformer/test_layer/test_layernorm.py
imports: from contextlib import nullcontext
  import torch
  import torch.nn
  from torch.testing import assert_close
  import colossalai
  from colossalai.lazy import LazyInitContext
  from colossalai.shardformer.layer import FusedLayerNorm
  from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
label_desc: Use image node_id 1 for calling a global function with example usage: check_layernorm(lazy_init) without return types