class_docstr: Context manager for lazy initialization. Enables initializing the model without allocating real memory.
  
  Args:
      tensor_cls (Union[_MyTensor, LazyTensor], optional): This is only for test. Defaults to LazyTensor.
      default_device (Optional[Union[torch.device, str, int]], optional): Defalt device for initialization.
          If it's cuda, initilization will be accelerated, but cuda memory will be allocated. By default, it's cpu.
          Defaults to None.
class_docstr_tok: ["Context","manager","for","lazy","initialization",".","Enables","initializing","the","model","without","allocating","real","memory",".","Args",":","tensor_cls","(","Union","[","_MyTensor",",","LazyTensor","]",",","optional",")",":","This","is","only","for","test",".","Defaults","to","LazyTensor",".","default_device","(","Optional","[","Union","[","torch.device",",","str",",","int","]","]",",","optional",")",":","Defalt","device","for","initialization",".","If","it","'s","cuda",",","initilization","will","be","accelerated",",","but","cuda","memory","will","be","allocated",".","By","default",",","it","'s","cpu",".","Defaults","to","None","."]
code: def __enter__(self):
      if LazyInitContext._replaced:
          raise RuntimeError(f"LazyInitContext is not reentrant")
      LazyInitContext._replaced = True
      self.old_default_device = self.tensor_cls.default_device
      self.tensor_cls.default_device = self.default_device
  
      def wrap_factory_method(target):
          # factory functions (eg. torch.empty())
          def wrapper(*args, **kwargs):
              return self.tensor_cls(target, *args, **kwargs)
  
          return wrapper, target
  
      def wrap_factory_like_method(orig_target, target):
          # factory_like functions (eg. torch.empty_like())
          def wrapper(*args, **kwargs):
              orig_t = args[0]
              return self.tensor_cls(
                  orig_target,
                  *orig_t.shape,
                  *args[1:],
                  device=orig_t.device,
                  dtype=orig_t.dtype,
                  **kwargs,
              )
  
          return wrapper, target
  
      def wrap_legacy_constructor(target, dtype):
          # legacy constructor (e.g. torch.LongTensor())
          def wrapper(*args, **kwargs):
              if len(args) == 1 and isinstance(args[0], torch.Tensor):
                  # (Tensor other)
                  return args[0]
              elif len(args) == 1:
                  # (object data, *, torch.device device)
                  kwargs = {**kwargs, "dtype": dtype}
                  replaced, orig = self.overrides["tensor"]
                  return replaced(*args, **kwargs)
              elif _is_int_tuple(args):
                  # (tuple of ints size, *, torch.device device)
                  kwargs = {**kwargs, "dtype": dtype}
                  replaced, orig = self.overrides["empty"]
                  return replaced(*args, **kwargs)
              else:
                  raise TypeError(
                      f"new() received an invalid combination of arguments - got {tuple(type(x) for x in args)}, but expected one of:\n * (Tensor other)\n * (tuple of ints size, *, torch.device device)\n * (object data, *, torch.device device)"
                  )
  
          return wrapper, target
  
      def wrap_no_meta_factory(target):
          # factory functions which don't support meta tensor backend
          def wrapper(*args, **kwargs):
              tensor = target(*args, **kwargs)
              return self.tensor_cls(lambda: None, concrete_data=tensor)
  
          return wrapper, target
  
      overrides = {
          target: wrap_factory_method(getattr(torch, target))
          for target in _NORMAL_FACTORY
          if callable(getattr(torch, target, None))
      }
  
      overrides.update(
          {
              target
              + "_like": wrap_factory_like_method(
                  getattr(torch, target),
                  getattr(torch, target + "_like"),
              )
              for target in _NORMAL_FACTORY
              if callable(getattr(torch, target + "_like", None))
          }
      )
  
      overrides.update(
          {
              target: wrap_legacy_constructor(
                  getattr(torch, target), dtype
              )
              for target, dtype in _LEGACY_TENSOR_CONSTRUCTOR.items()
              if callable(getattr(torch, target, None))
          }
      )
  
      overrides.update(
          {
              target: wrap_no_meta_factory(getattr(torch, target))
              for target in _NO_META_FACTORY
              if callable(getattr(torch, target, None))
          }
      )
  
      ConstructorManager.apply(overrides)
      PretrainedManager.inject()
  
imports: from types import MethodType
  from typing import Callable, Optional, Union
  import torch
  import torch.nn
  from packaging import version
  from torch import Tensor
  from torch.nn import Parameter
  from torch.utils._pytree import tree_map
  from colossalai.logging import get_dist_logger
  from .construction import ConstructorManager
  from .pretrained import PretrainedManager
  import colossalai._analyzer._subclasses._meta_registration
label_desc: Use image node_id 2 for calling the LazyInitContext obj's underlying member method code with example usage: obj.__enter__() and returns: wrapper,  target, wrapper,  target, wrapper,  target, wrapper,  target, self, self, self, args, replaced, replaced