code: def check_colo_module(
      module: torch.nn.Module, pg: ProcessGroup, recursive=True
  ):
      if is_colo_module(module):
          colo_module = get_colo_module(module)
          param_names = colo_module.get_param_names()
          compute_pattern = None
          for param_name in param_names:
              param = module.get_parameter(param_name)
              if not isinstance(param, ColoParameter):
                  raise Exception(
                      f"Invalid ColoParameter spec: {param} in {module} is not a ColoParameter."
                  )
              if param.has_compute_spec():
                  cur_compute_pattern = (
                      param.compute_spec.compute_pattern
                  )
                  if compute_pattern is None:
                      compute_pattern = cur_compute_pattern
                  else:
                      if cur_compute_pattern != compute_pattern:
                          raise Exception(
                              f"Invalid ColoParameter spec: Params in {module} have different compute_pattern."
                          )
              else:
                  continue
  
          if compute_pattern is not None:
              colo_module.register(compute_pattern, pg)
              if not colo_module.has_compute_pattern(compute_pattern):
                  raise Exception(
                      f"Invalid ColoParameter spec: ComputePattern {compute_pattern} in {module} is not allowed."
                  )
  
              match_specs = False
              allowed_specs = colo_module.get_dist_specs(
                  compute_pattern
              )
              for _, param_specs in allowed_specs.items():
                  cur_match = True
                  for param_name, dist_spec in param_specs.items():
                      param = module.get_parameter(param_name)
                      if param.has_compute_spec():
                          if dist_spec != param.dist_spec:
                              cur_match = False
                              break
                      else:
                          if dist_spec is not None:
                              cur_match = False
                              break
                  if cur_match == True:
                      match_specs = True
                      break
              if match_specs == False:
                  raise Exception(
                      f"Invalid ColoParameter spec: Params in {module} are incorrectly sharded."
                  )
      if recursive == True:
          for submodule in module.children():
              check_colo_module(submodule, pg=pg, recursive=True)
  
file_path: colossalai/colossalai/legacy/nn/parallel/layers/module_utils.py
imports: from typing import Dict
  import torch
  from colossalai.legacy.tensor import ComputeSpec, ProcessGroup
  from colossalai.tensor import ColoParameter
  from .None import ColoModule
label_desc: Use image node_id 4 for calling a global function with example usage: check_colo_module(module, pg, recursive) without return types