code: def named_params_with_colotensor(
      module: nn.Module,
      prefix: str = "",
      recurse: bool = True,
  ) -> Iterator[Tuple[str, Union[nn.Parameter, ColoTensor]]]:
      r"""Returns an iterator over module parameters (together with the
      ColoTensor parameters), yielding both the name of the parameter
      as well as the parameter itself. This is typically passed to a
      :class:torchshard._shard.sharded_optim.ShardedOptimizer
  
      Args:
          prefix (str): prefix to prepend to all parameter names.
          recurse (bool): if True, then yields parameters of this module
              and all submodules. Otherwise, yields only parameters that
              are direct members of this module.
  
      Yields:
          (string, Union[Tensor, ColoTensor]): Tuple containing
              the name and parameter (or ColoTensor parameter)
  
      Example:
  
          >>> model = torch.nn.Linear(*linear_size)
          >>> delattr(model.weight)
          >>> setattr(model.weight, ColoTensor(...))
          >>> for name, param in named_params_with_colotensor(model):
          >>>    if name in ['weight']:
          >>>        print(param.size())
  
      """
      modules = (
          module.named_modules(prefix=prefix)
          if recurse
          else [(prefix, module)]
      )
  
      memo = set()
      for mod_prefix, mod in modules:
          # find all sharded tensor params
          for name, val in vars(mod).items():
              if isinstance(val, ColoTensor) and val not in memo:
                  memo.add(val)
                  name = mod_prefix + ("." if mod_prefix else "") + name
                  yield name, val
  
      # find all nn.Parameters
      for name, val in module.named_parameters():
          yield name, val
  
file_path: colossalai/colossalai/tensor/utils.py
imports: from typing import Dict, Iterator, List, Tuple, Union
  import torch
  import torch.nn
  from colossalai.tensor.colo_tensor import ColoTensor
label_desc: Use image node_id 5 for calling a global function with example usage: named_params_with_colotensor(module, prefix, recurse) without return types