class_docstr: A naive implementation of LazyTensor (https://arxiv.org/pdf/2102.13267.pdf).
  
  Usage:
      1. Use ``LazyTensor`` instead of ``torch.Tensor``.
      >>> x = LazyTensor(torch.zeros, 2, 3)
      >>> x += 1
      >>> y = x * x
      >>> y = y.cuda().half()
      >>> y[0, 0] = 0
      >>> y = y.materialize()     # materialize the tensor
      >>> print(y)
      tensor([[0., 1., 1.],
              [1., 1., 1.]], device='cuda:0', dtype=torch.float16)
  
  Warnings:
      1. Cases that ``LazyTensor`` can't deal with.
      >>> x = LazyTensor(torch.ones, 2, 3)
      >>> x[0, 0] = -x[0, 0]    # this will cause infinite recursion
      >>> y = x.clone()
      >>> x.add_(1) # modifying origin tensor after cloning leads to wrong materialization
      >>> z = x.tolist()
      >>> x.zeros_() # modifying origin tensor after cloning tolist is not allowed
      >>> nn.utils.weight_norm(self.conv, name="weight", dim=2) # applying weight norm on a lazy tensor is not allowed
  
  
      2. Cases that ``LazyTensor`` becomes eager (early materialization).
      >>> b = a[:, 2:]  # get a slice of a lazy tensor triggers early materialization
      >>> chunks = a.split(3)  # this also triggers early materialization
      >>> x.data = torch.rand(2, 3) # directly setting data of a lazy tensor triggers early materialization
class_docstr_tok: ["A","naive","implementation","of","LazyTensor","(","https",":","\/\/arxiv.org\/pdf\/2102.13267.pdf",")",".","Usage",":","1",".","Use","``","LazyTensor","``","instead","of","``","torch.Tensor","``",".",">",">",">","x","=","LazyTensor","(","torch.zeros",",","2",",","3",")",">",">",">","x","+=","1",">",">",">","y","=","x","*","x",">",">",">","y","=","y.cuda","(",")",".half","(",")",">",">",">","y","[","0",",","0","]","=","0",">",">",">","y","=","y.materialize","(",")","#","materialize","the","tensor",">",">",">","print","(","y",")","tensor","(","[","[","0.",",","1.",",","1",".","]",",","[","1.",",","1.",",","1",".","]","]",",","device='cuda:0","'",",","dtype=torch.float16",")","Warnings",":","1",".","Cases","that","``","LazyTensor","``","ca","n't","deal","with",".",">",">",">","x","=","LazyTensor","(","torch.ones",",","2",",","3",")",">",">",">","x","[","0",",","0","]","=","-x","[","0",",","0","]","#","this","will","cause","infinite","recursion",">",">",">","y","=","x.clone","(",")",">",">",">","x.add_","(","1",")","#","modifying","origin","tensor","after","cloning","leads","to","wrong","materialization",">",">",">","z","=","x.tolist","(",")",">",">",">","x.zeros_","(",")","#","modifying","origin","tensor","after","cloning","tolist","is","not","allowed",">",">",">","nn.utils.weight_norm","(","self.conv",",","name=","''","weight","''",",","dim=2",")","#","applying","weight","norm","on","a","lazy","tensor","is","not","allowed","2",".","Cases","that","``","LazyTensor","``","becomes","eager","(","early","materialization",")",".",">",">",">","b","=","a","[",":",",","2",":","]","#","get","a","slice","of","a","lazy","tensor","triggers","early","materialization",">",">",">","chunks","=","a.split","(","3",")","#","this","also","triggers","early","materialization",">",">",">","x.data","=","torch.rand","(","2",",","3",")","#","directly","setting","data","of","a","lazy","tensor","triggers","early","materialization"]
code: def __torch_function__(cls, func, types, args=(), kwargs=None):
      if kwargs is None:
          kwargs = {}
      if func.__name__ in _EARLY_MATERIALIZED_OPS:
          # These OPs cannot be lazy and related tensors should be early materialized
          tree_map(cls._replace_with_materialized, args)
          tree_map(cls._replace_with_materialized, kwargs)
      is_inplace: bool = (
          func.__name__.endswith("_")
          and not (func.__name__.endswith("__"))
          or func.__name__ in ("__setitem__", "__set__")
      )
  
      is_change_meta_op: bool = func.__name__ in _CHANGE_META_OPS
  
      if isinstance(func, torch._C.ScriptMethod):
          # FIXME(ver217): torch script functions are not verified
  
          target = None
  
          def unwrap(x):
              if isinstance(x, LazyTensor):
                  return x._meta_data
              return x
  
          target: LazyTensor = args[0].clone()
          target._op_buffer.append((func, args, kwargs))
          target._meta_data = getattr(target._meta_data, func.name)(
              *tree_map(unwrap, args[1:]), **tree_map(unwrap, kwargs)
          )
          return target
      else:
          meta_to_lazy = {}
  
          def unwrap(x):
              if isinstance(x, LazyTensor):
                  if x._materialized_data is not None:
                      # for early materialized tensor, use its materialized data directly
                      return (
                          x._materialized_data
                          if is_change_meta_op
                          else x._materialized_data.data
                      )
                  t = x if is_inplace else x.clone()
                  if func.__name__ not in _NO_RERUN_OPS:
                      t._op_buffer.append((func, args, kwargs))
                  meta = (
                      x._meta_data
                      if is_change_meta_op
                      else x._meta_data.data
                  )
                  meta_to_lazy[meta] = t
                  return meta
              elif (
                  version.parse(torch.__version__)
                  >= version.parse("2.0.0")
                  and func.__name__ in _EXPAND_SCALAR_OPS
                  and not isinstance(x, torch.Tensor)
              ):
                  return _old_tensor_factory(x, device="meta")
              return x
  
          def wrap(y, i=None):
              if isinstance(y, torch.Tensor):
                  if y.is_meta:
                      if y in meta_to_lazy:
                          # inplace op, just return origin lazy tensor
                          return meta_to_lazy[y]
                      else:
                          # out of place op, create new lazy tensor
                          fn = (
                              lambda *a, **kw: func(*a, **kw)
                              if i is None
                              else func(*a, **kw)[i]
                          )
                          fn.__name__ = func.__name__
                          lazy_y = LazyTensor(
                              fn, *args, meta_data=y, **kwargs
                          )
                          return lazy_y
                  else:
                      # for early materialized tensor
                      return LazyTensor(lambda: None, concrete_data=y)
              return y
  
          cls._pre_op_fn()
          with ConstructorManager.disable():
              # to disable create lazy tensor in inner ops, this is a hack for torch 2.0
              o = func(
                  *tree_map(unwrap, args), **tree_map(unwrap, kwargs)
              )
          if isinstance(o, (tuple, list)):
              return type(o)(wrap(y, i=i) for i, y in enumerate(o))
          return wrap(o)
  
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 10 for calling the LazyTensor obj's underlying member method code with example usage: obj.__torch_function__(cls, func, types, args, kwargs) and returns: target, wrap, x, x, y, type, x, meta, unknown, _old_tensor_factory, LazyTensor, meta_to_lazy, lazy_y