code: def test_layernorm_meta_info(tensor_shape):
      meta_func = meta_register.get(torch.nn.LayerNorm)
  
      # construct input
      input_tensor = torch.rand(*tensor_shape, device="meta")
      output_tensor = torch.rand(*tensor_shape, device="meta")
      weight_tensor = torch.rand(tensor_shape[1], device="meta")
      bias_tensor = torch.rand(tensor_shape[1], device="meta")
  
      # construct operation data
      input_data = OperationData(
          name="input", type=OperationDataType.ARG, data=input_tensor
      )
  
      output_data = OperationData(
          name="output",
          type=OperationDataType.OUTPUT,
          data=output_tensor,
      )
  
      weight_data = OperationData(
          name="weight",
          type=OperationDataType.PARAM,
          data=weight_tensor,
      )
  
      bias_data = OperationData(
          name="bias", type=OperationDataType.PARAM, data=bias_tensor
      )
  
      # construct args and kwargs
      args = [input_data, output_data, weight_data, bias_data]
      kwargs = {"inplace": False}
  
      # estimated results
      (
          compute_cost,
          memory_cost,
          fwd_in,
          fwd_buffer,
          fwd_out,
      ) = meta_func(*args, **kwargs)
  
      # actual results
      input_real_tensor = torch.rand(*tensor_shape, device="cuda:0")
  
      input_real_tensor.requires_grad = True
  
      ln_module = torch.nn.LayerNorm(tensor_shape[1]).cuda()
  
      # fwd
      torch.cuda.reset_peak_memory_stats()
      mem_stamp0 = torch.cuda.memory_allocated()
      output_real_tensor = ln_module(input_real_tensor)
      fwd_allocated = torch.cuda.memory_allocated() - mem_stamp0
      fwd_peak = torch.cuda.max_memory_allocated() - mem_stamp0
  
      # bwd
      upstream_grad = torch.rand_like(output_real_tensor)
      torch.cuda.reset_peak_memory_stats()
      mem_stamp0 = torch.cuda.memory_allocated()
      torch.autograd.backward(output_real_tensor, upstream_grad)
      bwd_allocated = torch.cuda.memory_allocated() - mem_stamp0
      bwd_peak = torch.cuda.max_memory_allocated() - mem_stamp0
  
      compute_cost: TrainCycleItem
      memory_cost: TrainCycleItem
  
      print_results(
          [input_real_tensor],
          [output_real_tensor],
          compute_cost,
          memory_cost,
          fwd_allocated,
          fwd_peak,
          bwd_allocated,
          bwd_peak,
      )
  
file_path: ColossalAI/tests/test_auto_parallel/test_tensor_shard/test_metainfo/test_norm_metainfo.py
imports: import pytest
  import torch
  import torch.nn
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, OperationDataType, TrainCycleItem
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.initialize import launch
  from colossalai.logging import disable_existing_loggers
  from colossalai.testing.pytest_wrapper import run_on_environment_flag
  from colossalai.testing.utils import parameterize, rerun_if_address_is_in_use, spawn
  from tests.test_auto_parallel.test_tensor_shard.test_metainfo.utils import mem_test_for_node_strategy, print_results
label_desc: Use image node_id 3 for calling a global function with example usage: test_layernorm_meta_info(tensor_shape) without return types