code: def test_proxylessnas():
      net = ProxylessNAS()
      input = ShapeTensor(torch.randn(2, 3, 224, 224), True)
      output = shape_inference(net, input)
      assert output.real_shape == MutableShape(2, 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()
      }  # unwrap tuple
      output_shapes = {
          k: v[1].freeze(sample) for k, v in shapes.items()
      }
  
      expected_shapes = _real_shape_inference(
          real_model, torch.randn(2, 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()
      }
  
      left_keys = set(expected_input_shapes)
      for k in input_shapes:
          if k in expected_input_shapes and k in left_keys:
              assert input_shapes[k] == expected_input_shapes[k]
              assert output_shapes[k] == expected_output_shapes[k]
              left_keys.remove(k)
          elif k.startswith("blocks."):
              k_ = "blocks." + ".".join(
                  [d for d in k.split(".") if d.isdigit()]
              )
              if (
                  k_ in left_keys
                  and output_shapes[k] == expected_output_shapes[k_]
                  and input_shapes[k] == expected_input_shapes[k_]
              ):
                  left_keys.remove(k_)
      assert not left_keys
  
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 2 for calling a global function with example usage: test_proxylessnas() without return types