code: def _replace_module_with_type(
      root_module, init_fn, type_name, modules
  ):
      if modules is None:
          modules = []
  
      def apply(m):
          for name, child in m.named_children():
              if isinstance(child, type_name):
                  setattr(m, name, init_fn(child))
                  modules.append((child.key, getattr(m, name)))
              else:
                  apply(child)
  
      apply(root_module)
      return modules
  
imports: import logging
  from collections import OrderedDict
  import time
  import numpy
  import torch
  import nni.retiarii.nn.pytorch
  from nni.nas.pytorch.mutables import Mutable, InputChoice, LayerChoice
label_desc: Use image node_id 9 for calling a global function with example usage: _replace_module_with_type(root_module, init_fn, type_name, modules) and returns: modules