code: def test_repeat_blocks(model_cls):
      model = model_cls(4 * HIDDEN_DIM, HIDDEN_DIM, NUM_REPEAT_BLOCKS)
  
      tracer = ColoTracer(bias_addition_split=True)
      input_sample = {
          "x": torch.rand(BATCH_SIZE, SEQ_LENGTH, HIDDEN_DIM).to("meta")
      }
      graph = tracer.trace(root=model, meta_args=input_sample)
  
      gm = GraphModule(model, graph, model.__class__.__name__)
      shape_prop_pass(gm, *input_sample.values())
      gm.recompile()
  
      node_list = list(graph.nodes)
      root_module = graph.owning_module
      common_blocks = find_repeat_blocks(
          node_list, root_module, common_length_threshold=10
      )
  
      total_num_nodes = len(list(graph.nodes))
      # remove the input placeholder node and the output node
      num_repeat_nodes_per_block = (
          total_num_nodes - 2
      ) // NUM_REPEAT_BLOCKS
      for common_block in common_blocks:
          print(common_block)
      if model_cls == RepeatModel:
          assert len(common_blocks) == NUM_REPEAT_BLOCKS
          assert len(common_blocks[0]) == num_repeat_nodes_per_block
      elif model_cls == NonRepeatModel:
          assert len(common_blocks) == 0
  
imports: from typing import Optional, Tuple
  import torch
  import torch.nn
  from torch.fx import GraphModule
  from transformers.pytorch_utils import Conv1D
  from colossalai._analyzer.fx.passes import shape_prop_pass
  from colossalai._analyzer.fx.tracer.tracer import ColoTracer
  from colossalai.auto_parallel.tensor_shard.utils.factory import find_repeat_blocks
  from colossalai.testing import clear_cache_before_run, parameterize, run_on_environment_flag
label_desc: Use image node_id 1 for calling a global function with example usage: test_repeat_blocks(model_cls) without return types