code: def check_patch_embed():
      device = get_current_device()
      dtype = torch.float32
      j = gpc.get_local_rank(ParallelMode.PARALLEL_2D_ROW)
      i = gpc.get_local_rank(ParallelMode.PARALLEL_2D_COL)
  
      layer = PatchEmbedding2D(IMG_SIZE, 4, 3, HIDDEN_SIZE, dtype=dtype)
      torch.nn.init.ones_(layer.cls_token)
      torch.nn.init.ones_(layer.pos_embed)
      layer = layer.to(device)
  
      layer_master = VanillaPatchEmbedding(
          IMG_SIZE, 4, 3, HIDDEN_SIZE, dtype=dtype
      )
      torch.nn.init.ones_(layer_master.cls_token)
      torch.nn.init.ones_(layer_master.pos_embed)
      layer_master = layer_master.to(device)
  
      proj_weight_master = layer_master.weight.data
      torch.distributed.broadcast(proj_weight_master, src=0)
      proj_weight = torch.chunk(proj_weight_master, DEPTH, dim=0)[j]
      proj_weight = torch.chunk(proj_weight, DEPTH, dim=0)[i]
      layer.weight.data.copy_(proj_weight)
      proj_bias_master = layer_master.bias.data
      torch.distributed.broadcast(proj_bias_master, src=0)
      proj_bias = torch.chunk(proj_bias_master, DEPTH, dim=0)[j]
      proj_bias = torch.chunk(proj_bias, DEPTH, dim=0)[i]
      layer.bias.data.copy_(proj_bias)
  
      A_shape = (BATCH_SIZE, 3, IMG_SIZE, IMG_SIZE)
      A_master = torch.randn(A_shape, dtype=dtype, device=device)
      torch.distributed.broadcast(A_master, src=0)
      A = A_master.clone()
      out = layer(A)
  
      A_master = A_master.clone()
      C_master = layer_master(A_master)
      C = torch.chunk(C_master, DEPTH, dim=0)[i]
      C = torch.chunk(C, DEPTH, dim=-1)[j]
      check_equal(out, C)
      print_rank_0("patch embed forward: pass")
  
      grad_shape = C_master.shape
      grad_master = torch.randn(grad_shape, dtype=dtype, device=device)
      torch.distributed.broadcast(grad_master, src=0)
      grad = torch.chunk(grad_master, DEPTH, dim=0)[i]
      grad = torch.chunk(grad, DEPTH, dim=-1)[j]
      grad = grad.clone()
      out.backward(grad)
  
      grad_master = grad_master.clone()
      C_master.backward(grad_master)
  
      cls_grad_master = layer_master.cls_token.grad
      cls_grad = torch.chunk(cls_grad_master, DEPTH, dim=-1)[j]
      cls_grad = torch.chunk(cls_grad, DEPTH, dim=-1)[i]
      check_equal(cls_grad, layer.cls_token.grad)
  
      pos_grad_master = layer_master.pos_embed.grad
      pos_grad = torch.chunk(pos_grad_master, DEPTH, dim=-1)[j]
      pos_grad = torch.chunk(pos_grad, DEPTH, dim=-1)[i]
      check_equal(pos_grad, layer.pos_embed.grad)
  
      B_grad = layer_master.weight.grad
      B_grad = torch.chunk(B_grad, DEPTH, dim=0)[j]
      B_grad = torch.chunk(B_grad, DEPTH, dim=0)[i]
      check_equal(B_grad, layer.weight.grad)
  
      bias_grad = layer_master.bias.grad
      bias_grad = torch.chunk(bias_grad, DEPTH)[j]
      bias_grad = torch.chunk(bias_grad, DEPTH)[i]
      check_equal(bias_grad, layer.bias.grad)
      print_rank_0("patch embed backward: pass")
  
file_path: ColossalAI/tests/test_legacy/test_layers/test_2d/checks_2d/check_layer_2d.py
imports: import torch
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.nn import Classifier2D, CrossEntropyLoss2D, Embedding2D, LayerNorm2D, Linear2D, PatchEmbedding2D, VanillaClassifier, VanillaPatchEmbedding, VocabParallelClassifier2D, VocabParallelCrossEntropyLoss2D, VocabParallelEmbedding2D
  from colossalai.legacy.utils import print_rank_0
  from colossalai.utils import get_current_device
  from .common import BATCH_SIZE, DEPTH, HIDDEN_SIZE, IMG_SIZE, NUM_CLASSES, SEQ_LENGTH, VOCAB_SIZE, check_equal
label_desc: Use image node_id 4 for calling a global function with example usage: check_patch_embed() without return types