code: def test_wrapper_register():
      model = SimpleTorchModel()
  
      pruning_config_list = [
          {
              "op_names": ["conv1", "fc1"],
              "target_names": ["weight", "bias"],
              "sparse_ratio": 0.4,
          }
      ]
      wrappers, _ = register_wrappers(
          model, pruning_config_list, mode="pruning"
      )
      assert set(wrappers.keys()) == set(["conv1", "fc1"])
  
      quantization_config_list = [
          {
              "op_names": ["conv2", "fc1"],
              "target_names": ["_input_", "weight", "_output_"],
              "quant_dtype": "int8",
              "quant_scheme": "affine",
          }
      ]
      wrappers, _ = register_wrappers(
          model,
          quantization_config_list,
          mode="quantization",
          existed_wrappers=wrappers,
      )
      assert set(wrappers.keys()) == set(["conv1", "conv2", "fc1"])
  
      distillation_config_list = [
          {
              "op_names": ["conv3", "fc1"],
              "target_names": ["_output_"],
              "lambda": 0.1,
          }
      ]
      wrappers, _ = register_wrappers(
          model,
          distillation_config_list,
          mode="distillation",
          existed_wrappers=wrappers,
      )
      assert set(wrappers.keys()) == set(
          ["conv1", "conv2", "conv3", "fc1"]
      )
  
file_path: nni/test/algo/compression/base/test_wrapper.py
imports: from __future__ import annotations
  import pytest
  import torch
  from torch.nn.functional import conv2d
  from nni.compression.base.wrapper import ModuleWrapper, register_wrappers
  from ..assets.simple_mnist import SimpleTorchModel
label_desc: Use image node_id 2 for calling a global function with example usage: test_wrapper_register() without return types