code: def test_shufflenet():
      net = ShuffleNetSpace(channel_search=True)
      input = ShapeTensor(torch.randn(1, 3, 224, 224), True)
      output = shape_inference(net, input)
      shapes = submodule_input_output_shapes(net, input)
      assert output.real_shape == MutableShape(1, 1000)
  
      sample = {}
      real_model = net.random(sample)
      shapes = submodule_input_output_shapes(net, input)
      input_shapes = {
          k: v[0][0].freeze(sample) for k, v in shapes.items()
      }
      output_shapes = {
          k: v[1].freeze(sample) for k, v in shapes.items()
      }
  
      expected_shapes = _real_shape_inference(
          real_model, torch.randn(1, 3, 224, 224)
      )
      expected_input_shapes = {
          k: v[0][0] for k, v in expected_shapes.items()
      }
      expected_output_shapes = {
          k: v[1] for k, v in expected_shapes.items()
      }
  
      for k in expected_input_shapes:
          if k in input_shapes:
              assert input_shapes[k] == expected_input_shapes[k]
              assert output_shapes[k] == expected_output_shapes[k]
          else:
              kpre, kpre2, kpost = k.split(
                  ".", 2
              )  # features 0 branch_main.1
              kpre = kpre + "." + kpre2
              for k_ in input_shapes:
                  if (
                      k_.startswith(kpre)
                      and k_.endswith(kpost)
                      and input_shapes[k_] == expected_input_shapes[k]
                      and output_shapes[k_] == expected_output_shapes[k]
                  ):
                      break
              else:
                  assert False, f"{k} not found"
  
file_path: nni/test/algo/nas/profiler/test_shape.py
imports: import re
  from functools import partial
  import torch
  from nni.nas.hub.pytorch import *
  from nni.nas.profiler.pytorch.utils import *
label_desc: Use image node_id 5 for calling a global function with example usage: test_shufflenet() without return types