code: def model_space_init_wrapper(
      original_init_fn: Callable[..., None]
  ) -> Callable[..., None]:
      """Wrap the ``__init__`` inside a model namespace.
  
      This should only be wrapped on *subclasses* of :class:`ModelSpace`.
      It could possibly be wrapped multiple times when subclass of :class:`ModelSpace` is further subclassed.
      Currently, nested label namespace and dry run context will be created.
      """
  
      def init_with_context(self: ModelSpace, *args, **kwargs):
          arch = current_model()
          if arch is None:
              # If no arch is set, we are in dry run mode.
              # This is the only place where we set the dry run mode.
              # The context will record the sample generated during ``__init__``.
              # Some mutables use :func:`ensure_frozen` to freeze the sample.
              # The context is to make sure same label always return the same choice.
              self._frozen_context = frozen_context()
              with self._frozen_context:
                  return original_init_fn(self, *args, **kwargs)
          else:
              # Already in context.
              self._frozen_context = frozen_context.top_context()
              return original_init_fn(self, *args, **kwargs)
  
      def new_init(self: ModelSpace, *args, **kwargs) -> None:
          # Save init arguments first. Skips if already saved.
          self.auto_save_init_arguments(*args, **kwargs)
  
          if not hasattr(self, "_label_scope"):
              if self._label_prefix is not None:
                  self._label_scope = label_scope(self._label_prefix)
              else:
                  self._label_scope = strict_label_scope(
                      "_unused_"
                  )  # the name is not used
          if hasattr(self, "_label_scope") and not self._label_scope.activated:  # type: ignore
              # Has a label scope but it's not activated. Create a "with".
              with self._label_scope:  # type: ignore
                  return init_with_context(self, *args, **kwargs)
          else:
              return init_with_context(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 2 for calling a global function with example usage: model_space_init_wrapper(original_init_fn) and returns: new_init, original_init_fn, init_with_context, original_init_fn, init_with_context