code: def test_comm_size_compute():
      model = MLP(MODEL_DIM)
      input_sample = torch.rand(BATCH_SIZE, MODEL_DIM, device="meta")
      gm = symbolic_trace(model)
      if is_compatible:
          input_sample = MetaTensor(
              input_sample, fake_device=next(gm.parameters()).device
          )
      MetaInfoProp(gm).run(input_sample)
      annotated_model = uniform_split_pass(gm, PIPELINE_SIZE)
      split_model, split_submodules = split_with_split_nodes_pass(
          annotated_model
      )
      submodule_list = list(split_model.children())
      comm_size = get_comm_size(submodule_list[0], submodule_list[1])
      # the shape of tensor send from partition 0 to partition 1 is (8, 16)
      assert comm_size == 128
  
file_path: colossalai/tests/test_fx/test_comm_size_compute.py
imports: import torch
  from torch.fx import symbolic_trace
  from colossalai.fx._compatibility import is_compatible_with_meta
  from colossalai.fx.passes.adding_split_node_pass import split_with_split_nodes_pass, uniform_split_pass
  from colossalai.fx.passes.meta_info_prop import MetaInfoProp
  from colossalai.fx.passes.utils import get_comm_size
  from colossalai.testing import clear_cache_before_run
label_desc: Use image node_id 1 for calling a global function with example usage: test_comm_size_compute() without return types