code: def parametrized_module_init_wrapper(
      original_init_fn: Callable[..., None]
  ) -> Callable[..., None]:
      """Wrap the ``__init__`` for :class:`ParametrizedModule`.
  
      It returns a function, which first records all the parameters passed to ``__init__``,
      then ensures all the parameters are frozen, and finally calls the original ``__init__``.
      """
  
      def new_init(self: ParametrizedModule, *args, **kwargs):
          arch = current_model()
          if arch is not None:
              # This is the case where fixed-arch context exists,
              # yet ParametrizedModule doesn't have a bound type.
              args, kwargs = self.freeze_init_arguments(
                  arch, *args, **kwargs
              )
              return original_init_fn(self, *args, **kwargs)
  
          # This is the case where fixed-arch context doesn't exist.
          self.auto_save_init_arguments(*args, **kwargs)
          for arg in itertools.chain(args, kwargs.values()):
              if isinstance(arg, Mutable):
                  self.add_mutable(arg)
              else:
                  _warn_if_nested_mutable(arg, self.__class__.__name__)
          # Sometimes, arguments will be hijacked to make the inner wrapped class happy.
          # For example Conv2d(choice([3, 5, 7])) should be Conv2d(3) instead,
          # because Conv2d doesn't recognize choice([3, 5, 7]).
          args, kwargs = self.freeze_init_arguments(
              None, *args, **kwargs
          )
          return original_init_fn(self, *args, **kwargs)
  
      return new_init
  
imports: import copy
  import itertools
  import inspect
  import logging
  from typing import Any, Callable, Type, Iterable, Dict, Tuple, List, Optional, TypeVar
  import torch
  from torch import nn
  from nni.mutable import Mutable, LabeledMutable, Sample, SampleValidationError, ensure_frozen, frozen_context, label_scope, frozen
  from nni.nas.space import BaseModelSpace, current_model, model_context
  from nni.common.serializer import SerializableObject, _formulate_arguments, _copy_class_wrapper_attributes
label_desc: Use image node_id 3 for calling a global function with example usage: parametrized_module_init_wrapper(original_init_fn) and returns: new_init, original_init_fn, original_init_fn