code: def _assert_output_shape(
      data, module, patch_fn, expect_exception, output_shape
  ):
      output = _run(data, module, patch_fn)
  
      if expect_exception:
          assert isinstance(output, AssertionError)
      else:
          assert not isinstance(output, Exception)
          if isinstance(output, tuple):
              for item, shape in zip(output, output_shape):
                  assert item.is_meta
                  assert item.shape == shape
          else:
              assert output.is_meta
              assert output.shape == output_shape
  
file_path: colossalai/tests/test_fx/test_tracer/test_patched_module.py
imports: import torch
  from colossalai.fx.tracer.meta_patch import patched_module
  from colossalai.testing import clear_cache_before_run
label_desc: Use image node_id 2 for calling a global function with example usage: _assert_output_shape(data, module, patch_fn, expect_exception, output_shape) without return types