code: def check_embedding_1d(lazy_init: bool):
      ctx = LazyInitContext() if lazy_init else nullcontext()
  
      embedding = nn.Embedding(32, 128).cuda()
      with ctx:
          embedding_copy = nn.Embedding(32, 128).cuda()
      embedding_1d = Embedding1D.from_native_module(
          embedding_copy, process_group=None
      )
  
      assert embedding_1d.weight.shape == torch.Size([32, 64])
      assert embedding_1d.weight is embedding_copy.weight
  
      # ensure state dict is reversibly loadable
      embedding.load_state_dict(embedding_1d.state_dict())
      embedding_1d.load_state_dict(embedding.state_dict())
  
      # check computation correctness
      x = torch.randint(low=0, high=32, size=(4, 32)).cuda()
      out = embedding(x)
      gather_out = embedding_1d(x)
      assert_close(out, gather_out)
  
      # check backward correctness
      out.sum().backward()
      gather_out.sum().backward()
  
      rank = dist.get_rank()
      target_grad = torch.chunk(embedding.weight.grad, 2, dim=1)[rank]
      assert_close(target_grad, embedding_1d.weight.grad)
  
file_path: colossalai/tests/test_shardformer/test_layer/test_embedding.py
imports: from contextlib import nullcontext
  import torch
  import torch.distributed
  import torch.nn
  from torch.testing import assert_close
  import colossalai
  from colossalai.lazy import LazyInitContext
  from colossalai.shardformer.layer import Embedding1D
  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_embedding_1d(lazy_init) without return types