_tensor.py 75 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893
  1. # mypy: allow-untyped-defs
  2. import copyreg
  3. import enum
  4. import functools
  5. import itertools
  6. import warnings
  7. from collections import OrderedDict
  8. from collections.abc import Callable
  9. from copy import deepcopy
  10. from numbers import Number
  11. from typing import Any, cast, Concatenate, TypeVar, Union
  12. from typing_extensions import ParamSpec
  13. import torch
  14. import torch._C as _C
  15. from torch._namedtensor_internals import (
  16. check_serializing_named_tensor,
  17. is_ellipsis,
  18. resolve_ellipsis,
  19. single_ellipsis_index,
  20. unzip_namedshape,
  21. update_names,
  22. )
  23. from torch.overrides import (
  24. get_default_nowrap_functions,
  25. handle_torch_function,
  26. has_torch_function,
  27. has_torch_function_unary,
  28. has_torch_function_variadic,
  29. )
  30. _P = ParamSpec("_P")
  31. _TensorLike = TypeVar("_TensorLike", bound=_C.TensorBase)
  32. def _handle_torch_function_and_wrap_type_error_to_not_implemented(
  33. f: Callable[Concatenate[_TensorLike, _P], "Tensor"],
  34. ) -> Callable[Concatenate[_TensorLike, _P], "Tensor"]:
  35. @functools.wraps(f)
  36. def wrapped(self: _TensorLike, *args: _P.args, **kwargs: _P.kwargs) -> "Tensor":
  37. try:
  38. # See https://github.com/pytorch/pytorch/issues/75462
  39. sargs = self, *args
  40. if has_torch_function(sargs):
  41. return handle_torch_function(wrapped, sargs, *sargs, **kwargs)
  42. return f(self, *args, **kwargs)
  43. except TypeError:
  44. return NotImplemented
  45. return wrapped
  46. # Should not be used, this is kept only for BC of loading old serialized Tensor subclasses
  47. def _rebuild_from_type(func, type, args, dict):
  48. if type is Tensor:
  49. return func(*args)
  50. ret = func(*args).as_subclass(type)
  51. ret.__dict__ = dict
  52. return ret
  53. def _rebuild_from_type_v2(func, new_type, args, state):
  54. ret = func(*args)
  55. if type(ret) is not new_type:
  56. ret = ret.as_subclass(new_type)
  57. # Tensor does define __setstate__ even though it doesn't define
  58. # __getstate__. So only use __setstate__ if it is NOT the one defined
  59. # on Tensor
  60. if (
  61. getattr(ret.__class__, "__setstate__", Tensor.__setstate__)
  62. is not Tensor.__setstate__
  63. ):
  64. ret.__setstate__(state)
  65. else:
  66. ret = torch._utils._set_obj_state(ret, state)
  67. return ret
  68. def _dtype_to_typestr(dtype):
  69. # CUDA devices are little-endian and tensors are stored in native byte
  70. # order. 1-byte entries are endian-agnostic.
  71. return {
  72. torch.complex64: "<c8",
  73. torch.complex128: "<c16",
  74. torch.bfloat16: "<V2", # Same as ml_dtypes.bfloat16.dtype.str.
  75. torch.float16: "<f2",
  76. torch.float32: "<f4",
  77. torch.float64: "<f8",
  78. torch.uint8: "|u1",
  79. torch.int8: "|i1",
  80. torch.uint16: "<u2",
  81. torch.int16: "<i2",
  82. torch.uint32: "<u4",
  83. torch.int32: "<i4",
  84. torch.uint64: "<u8",
  85. torch.int64: "<i8",
  86. torch.bool: "|b1",
  87. }[dtype]
  88. # NB: If you subclass Tensor, and want to share the subclassed class
  89. # across processes, you must also update torch/multiprocessing/reductions.py
  90. # to define a ForkingPickler serialization mode for the class.
  91. #
  92. # NB: If you add a new method to Tensor, you must update
  93. # torch/_C/__init__.pyi.in to add a type annotation for your method;
  94. # otherwise, it will not show up in autocomplete.
  95. class Tensor(torch._C.TensorBase):
  96. _is_param: bool
  97. # pyrefly: ignore [missing-attribute]
  98. __dlpack_c_exchange_api__: object = torch._C._dlpack_exchange_api()
  99. def _clear_non_serializable_cached_data(self):
  100. r"""Clears any data cached in the tensor's ``__dict__`` that would prevent the tensor
  101. from being serialized.
  102. For example, subclasses with custom dispatched sizes / strides cache this info in
  103. non-serializable PyCapsules within the ``__dict__``, and this must be cleared out for
  104. serialization to function.
  105. Any subclass that overrides this MUST call ``super()._clear_non_serializable_cached_data().``
  106. Additional data cleared within the override must be able to be re-cached transparently
  107. to avoid breaking subclass functionality.
  108. """
  109. if has_torch_function_unary(self):
  110. return handle_torch_function(
  111. Tensor._clear_non_serializable_cached_data, (self,), self
  112. )
  113. # NB: Wrapper subclasses that implement custom-dispatched sizes / strides cache
  114. # this info via non-serializable PyCapsules.
  115. CACHED_SIZES_STRIDES_KEYS = [
  116. "_sym_sizes_capsule",
  117. "_sym_sizes_capsule_len",
  118. "_sym_strides_capsule",
  119. "_sym_strides_capsule_len",
  120. ]
  121. for key in CACHED_SIZES_STRIDES_KEYS:
  122. self.__dict__.pop(key, None)
  123. def __deepcopy__(self, memo):
  124. if has_torch_function_unary(self):
  125. return handle_torch_function(Tensor.__deepcopy__, (self,), self, memo)
  126. if not self.is_leaf:
  127. raise RuntimeError(
  128. "Only Tensors created explicitly by the user "
  129. "(graph leaves) support the deepcopy protocol at the moment. "
  130. "If you were attempting to deepcopy a module, this may be because "
  131. "of a torch.nn.utils.weight_norm usage, "
  132. "see https://github.com/pytorch/pytorch/pull/103001"
  133. )
  134. if id(self) in memo:
  135. return memo[id(self)]
  136. with torch.no_grad():
  137. # TODO: skipping storage copy is wrong for meta, as meta
  138. # does accurate alias tracking; however, the code below
  139. # doesn't work because of
  140. # https://github.com/pytorch/pytorch/issues/47442
  141. # Update the test in test_serialization if you remove 'meta' from here
  142. if (
  143. self.is_sparse
  144. or self.device.type
  145. in ["lazy", "xla", "mtia", "mps", "maia", "meta", "ipu"]
  146. or (
  147. not torch._C._has_storage(self)
  148. and self.device.type == torch._C._get_privateuse1_backend_name()
  149. )
  150. or (type(self) is not Tensor and self.data_ptr() == 0)
  151. ):
  152. new_tensor = self.clone()
  153. if type(new_tensor) is not type(self):
  154. raise RuntimeError(
  155. "The default implementation of __deepcopy__() for wrapper subclasses "
  156. "only works for subclass types that implement clone() and for which "
  157. "cloning returns another instance of the same subclass. You should either "
  158. "properly implement clone() for your subclass or override __deepcopy__() "
  159. "if it is intended behavior for clone() to return an instance of a "
  160. "different type."
  161. )
  162. else:
  163. new_storage = self._typed_storage()._deepcopy(memo)
  164. if self.is_quantized:
  165. # quantizer_params can be different type based on torch attribute
  166. quantizer_params: (
  167. tuple[torch.qscheme, float, int]
  168. | tuple[torch.qscheme, Tensor, Tensor, int]
  169. )
  170. if self.qscheme() == torch.per_tensor_affine:
  171. quantizer_params = (
  172. self.qscheme(),
  173. self.q_scale(),
  174. self.q_zero_point(),
  175. )
  176. elif self.qscheme() in (
  177. torch.per_channel_affine,
  178. torch.per_channel_affine_float_qparams,
  179. ):
  180. quantizer_params = (
  181. self.qscheme(),
  182. self.q_per_channel_scales(),
  183. self.q_per_channel_zero_points(),
  184. self.q_per_channel_axis(),
  185. )
  186. else:
  187. raise RuntimeError(
  188. f"Unsupported qscheme {self.qscheme()} in deepcopy"
  189. )
  190. # TODO: Once we decide to break serialization FC, no longer
  191. # need to wrap with TypedStorage
  192. new_tensor = torch._utils._rebuild_qtensor(
  193. torch.storage.TypedStorage(
  194. wrap_storage=new_storage._untyped_storage,
  195. dtype=self.dtype,
  196. _internal=True,
  197. ),
  198. self.storage_offset(),
  199. self.size(),
  200. self.stride(),
  201. quantizer_params,
  202. self.requires_grad,
  203. self._backward_hooks,
  204. )
  205. if type(new_tensor) is not type(self):
  206. raise RuntimeError(
  207. "The default implementation of __deepcopy__() for quantized tensors "
  208. "expects the tensor returned by torch._utils._rebuild_qtensor() to "
  209. "match the type of the instance being copied. If you encounter this, "
  210. "please open an issue on PyTorch's GitHub."
  211. )
  212. else:
  213. new_tensor = self.new_empty([])
  214. if type(new_tensor) is not type(self):
  215. raise RuntimeError(
  216. "The default implementation of __deepcopy__() for non-wrapper subclasses "
  217. "only works for subclass types that implement new_empty() and for which "
  218. "that function returns another instance of the same subclass. You should "
  219. "either properly implement new_empty() for your subclass or override "
  220. "__deepcopy__() if it is intended behavior for new_empty() to return "
  221. "an instance of a different type."
  222. )
  223. new_tensor.set_(
  224. new_storage, self.storage_offset(), self.size(), self.stride()
  225. )
  226. if self.is_conj():
  227. new_tensor = new_tensor.conj_physical()
  228. if self.is_neg():
  229. new_tensor = new_tensor.neg()
  230. if self.requires_grad:
  231. new_tensor.requires_grad_()
  232. if self.grad is not None:
  233. new_tensor.grad = self.grad.__deepcopy__(memo)
  234. if type(self) is not Tensor:
  235. if type(new_tensor) is not type(self):
  236. raise RuntimeError(
  237. "Type of deepcopy result does not match the type of the source tensor. "
  238. "If you encounter this, please open an issue on PyTorch's GitHub."
  239. )
  240. # Plain Tensors don't have slots
  241. slots_to_save = copyreg._slotnames(self.__class__) # type: ignore[attr-defined]
  242. for slot in slots_to_save:
  243. if hasattr(self, slot):
  244. setattr(new_tensor, slot, deepcopy(getattr(self, slot), memo))
  245. # don't try to deepcopy non-serializable cached data
  246. self._clear_non_serializable_cached_data()
  247. new_tensor.__dict__ = deepcopy(self.__dict__, memo)
  248. memo[id(self)] = new_tensor
  249. return new_tensor
  250. def __reduce_ex__(self, proto):
  251. materialize_fake_tensors = (
  252. torch.serialization._serialization_tls.materialize_fake_tensors
  253. )
  254. state = torch._utils._get_obj_state(self)
  255. # Ignore all state when using FakeTensor with skip_data(materialize_fake_tensors) because FakeTensor has
  256. # some state that cannot be pickled
  257. if (
  258. # TODO: remove hasattr, it's a hack to support versions of torch that
  259. # don't have _subclasses
  260. hasattr(torch, "_subclasses")
  261. and type(self) is torch._subclasses.fake_tensor.FakeTensor
  262. and materialize_fake_tensors
  263. ) or (type(self) is Tensor and not state):
  264. # Fast path for regular tensor without Python state.
  265. return self._reduce_ex_internal(proto)
  266. if has_torch_function_unary(self):
  267. return handle_torch_function(Tensor.__reduce_ex__, (self,), self, proto)
  268. func, args = self._reduce_ex_internal(proto)
  269. # sizes / strides cache needs to be cleared here because it'll just be re-cached
  270. # if cleared earlier. Note that state references the -actual- tensor dict.
  271. self._clear_non_serializable_cached_data()
  272. return (_rebuild_from_type_v2, (func, type(self), args, state))
  273. def storage(self):
  274. r"""
  275. storage() -> torch.TypedStorage
  276. Returns the underlying :class:`TypedStorage`.
  277. .. warning::
  278. :class:`TypedStorage` is deprecated. It will be removed in the future, and
  279. :class:`UntypedStorage` will be the only storage class. To access the
  280. :class:`UntypedStorage` directly, use :attr:`Tensor.untyped_storage()`.
  281. """
  282. if has_torch_function_unary(self):
  283. return handle_torch_function(Tensor.storage, (self,), self)
  284. torch.storage._warn_typed_storage_removal(stacklevel=2)
  285. return self._typed_storage()
  286. # For internal use only, to avoid raising deprecation warning
  287. def _typed_storage(self):
  288. untyped_storage = self.untyped_storage()
  289. return torch.TypedStorage(
  290. wrap_storage=untyped_storage, dtype=self.dtype, _internal=True
  291. )
  292. def _reduce_ex_internal(self, proto):
  293. check_serializing_named_tensor(self)
  294. from torch.utils.hooks import warn_if_has_hooks
  295. # See Note [Don't serialize hooks]
  296. warn_if_has_hooks(self)
  297. backward_hooks: dict[Any, Any] = OrderedDict()
  298. skip_data = torch.serialization._serialization_tls.skip_data
  299. materialize_fake_tensors = (
  300. torch.serialization._serialization_tls.materialize_fake_tensors
  301. )
  302. if self.device.type in ["xla", "maia", "mtia"] or (
  303. not torch._C._has_storage(self)
  304. and self.device.type == torch._C._get_privateuse1_backend_name()
  305. ):
  306. if skip_data:
  307. raise RuntimeError(
  308. "Cannot serialize tensors on backends with no storage under skip_data context manager"
  309. )
  310. cpu_tensor = self.cpu()
  311. return (
  312. torch._utils._rebuild_device_tensor_from_cpu_tensor,
  313. (cpu_tensor, self.dtype, str(self.device), self.requires_grad),
  314. )
  315. if self.device.type == "meta":
  316. # NB: This implementation BREAKS storage sharing. Current
  317. # hypothesis is that no one cares for meta tensors.
  318. if skip_data:
  319. warnings.warn(
  320. "Serializing tensors on the meta device under skip_data context manager is a no-op",
  321. stacklevel=2,
  322. )
  323. arg_meta = (
  324. self.dtype,
  325. tuple(self.size()),
  326. self.stride(),
  327. self.requires_grad,
  328. )
  329. return (torch._utils._rebuild_meta_tensor_no_storage, arg_meta)
  330. if self.is_quantized:
  331. if skip_data:
  332. raise RuntimeError(
  333. "Cannot serialize qtensor under skip_data context manager, file an issue if you need this feature"
  334. )
  335. # quantizer_params can be different type based on torch attribute
  336. quantizer_params: (
  337. tuple[torch.qscheme, float, int] | tuple[Any, Tensor, Tensor, int]
  338. )
  339. if self.qscheme() == torch.per_tensor_affine:
  340. quantizer_params = (
  341. torch.per_tensor_affine,
  342. self.q_scale(),
  343. self.q_zero_point(),
  344. )
  345. elif self.qscheme() in (
  346. torch.per_channel_affine,
  347. torch.per_channel_affine_float_qparams,
  348. ):
  349. # convert scales and zero points to tuple to avoid recursive calls
  350. # when/if we get multi-axis quantized tensors in the future, the shape
  351. # is recoverable from the main tensor shape
  352. quantizer_params = (
  353. torch.per_channel_affine,
  354. self.q_per_channel_scales(),
  355. self.q_per_channel_zero_points(),
  356. self.q_per_channel_axis(),
  357. )
  358. else:
  359. raise RuntimeError(
  360. f"Serialization is not supported for tensors of type {self.qscheme()}"
  361. )
  362. # TODO: Once we decide to break serialization FC, no longer
  363. # need to wrap with TypedStorage
  364. args_qtensor = (
  365. torch.storage.TypedStorage(
  366. wrap_storage=self._typed_storage()._untyped_storage,
  367. dtype=self.dtype,
  368. _internal=True,
  369. ),
  370. self.storage_offset(),
  371. tuple(self.size()),
  372. self.stride(),
  373. quantizer_params,
  374. self.requires_grad,
  375. backward_hooks,
  376. )
  377. return (torch._utils._rebuild_qtensor, args_qtensor)
  378. elif self.is_sparse:
  379. if self.layout == torch.sparse_coo:
  380. args_sparse = (
  381. self.layout,
  382. (self._indices(), self._values(), self.size(), self.is_coalesced()),
  383. )
  384. else:
  385. raise NotImplementedError(
  386. f"sparse tensor __reduce_ex__ for layout `{self.layout}`"
  387. )
  388. return (torch._utils._rebuild_sparse_tensor, args_sparse)
  389. elif self.layout in {
  390. torch.sparse_csr,
  391. torch.sparse_csc,
  392. torch.sparse_bsr,
  393. torch.sparse_bsc,
  394. }:
  395. if self.layout in {torch.sparse_csr, torch.sparse_bsr}:
  396. compressed_indices, plain_indices = (
  397. self.crow_indices(),
  398. self.col_indices(),
  399. )
  400. else:
  401. compressed_indices, plain_indices = (
  402. self.ccol_indices(),
  403. self.row_indices(),
  404. )
  405. args_sparse_compressed = (
  406. self.layout,
  407. (
  408. compressed_indices,
  409. plain_indices,
  410. self.values(),
  411. self.size(),
  412. ),
  413. )
  414. return (torch._utils._rebuild_sparse_tensor, args_sparse_compressed)
  415. elif self.is_nested:
  416. if skip_data:
  417. raise RuntimeError(
  418. "Cannot serialize nested tensor under skip_data context manager, file an issue if you need this feature"
  419. )
  420. args_nested = (
  421. # NB: values() currently returns the storage as a buffer in an unsafe way.
  422. # Ideally, we'd use a private API for this instead. TODO: Switch to this if
  423. # we ever get around to adding it.
  424. self.values(),
  425. self._nested_tensor_size(),
  426. self._nested_tensor_strides(),
  427. self._nested_tensor_storage_offsets(),
  428. )
  429. return (torch._utils._rebuild_nested_tensor, args_nested)
  430. elif (
  431. type(self) is not torch.Tensor
  432. and type(self).__torch_dispatch__ is not torch.Tensor.__torch_dispatch__
  433. and (
  434. isinstance(self, torch._subclasses.functional_tensor.FunctionalTensor)
  435. or (
  436. not isinstance(self, torch._subclasses.fake_tensor.FakeTensor)
  437. and self.data_ptr() == 0
  438. )
  439. )
  440. ):
  441. arg_wrapper_subclass = (
  442. type(self),
  443. self.dtype,
  444. tuple(self.size()),
  445. self.stride(),
  446. self.storage_offset(),
  447. self.layout,
  448. self.device,
  449. self.requires_grad,
  450. )
  451. return (torch._utils._rebuild_wrapper_subclass, arg_wrapper_subclass)
  452. elif (
  453. type(self) is not torch.Tensor
  454. and type(self).__torch_dispatch__ is not torch.Tensor.__torch_dispatch__
  455. and (
  456. isinstance(self, torch._subclasses.fake_tensor.FakeTensor)
  457. and not (skip_data and materialize_fake_tensors)
  458. )
  459. ):
  460. arg_wrapper_subclass = (
  461. type(self),
  462. self.dtype,
  463. tuple(self.size()),
  464. self.stride(),
  465. self.storage_offset(),
  466. self.layout,
  467. self.device,
  468. self.requires_grad,
  469. )
  470. return (torch._utils._rebuild_wrapper_subclass, arg_wrapper_subclass)
  471. else:
  472. v3_dtypes = torch.storage._new_dtypes()
  473. if self.dtype in v3_dtypes:
  474. rebuild_func = torch._utils._rebuild_tensor_v3
  475. storage = self.untyped_storage()
  476. else:
  477. # TODO: Once we decide to break serialization FC, no longer
  478. # need to wrap with TypedStorage
  479. rebuild_func = torch._utils._rebuild_tensor_v2 # type: ignore[assignment]
  480. storage = torch.storage.TypedStorage(
  481. wrap_storage=self._typed_storage()._untyped_storage,
  482. dtype=self.dtype,
  483. _internal=True,
  484. ) # type: ignore[assignment]
  485. # TODO: remove hasattr, it's a hack to support versions of torch that
  486. # don't have _subclasses
  487. if (
  488. hasattr(torch, "_subclasses")
  489. and isinstance(self, torch._subclasses.fake_tensor.FakeTensor)
  490. and skip_data
  491. ):
  492. storage._fake_device = self.device
  493. args = (
  494. storage,
  495. self.storage_offset(),
  496. tuple(self.size()),
  497. self.stride(),
  498. self.requires_grad,
  499. backward_hooks,
  500. ) # previously was self._backward_hooks
  501. if isinstance(storage, torch.storage.UntypedStorage):
  502. args = args + (self.dtype,) # type: ignore[assignment]
  503. metadata = torch._utils.get_tensor_metadata(self)
  504. if metadata:
  505. args = args + (metadata,) # type: ignore[assignment]
  506. return (rebuild_func, args)
  507. def __setstate__(self, state):
  508. if has_torch_function_unary(self):
  509. return handle_torch_function(Tensor.__setstate__, (self,), self, state)
  510. # Warning: this method is NOT called when you torch.load() a tensor;
  511. # that is managed by _rebuild_tensor_v2
  512. if not self.is_leaf:
  513. raise RuntimeError("__setstate__ can be only called on leaf Tensors")
  514. if len(state) == 4:
  515. # legacy serialization of Tensor
  516. self.set_(*state)
  517. return
  518. elif len(state) == 5:
  519. # legacy serialization of Variable
  520. self.data = state[0]
  521. state = (state[3], state[4], state[2])
  522. # The setting of _backward_hooks is expected to be a no-op.
  523. # See Note [Don't serialize hooks]
  524. self.requires_grad, _, self._backward_hooks = state
  525. def __repr__(self, *, tensor_contents=None):
  526. if has_torch_function_unary(self):
  527. return handle_torch_function(
  528. Tensor.__repr__, (self,), self, tensor_contents=tensor_contents
  529. )
  530. # All strings are unicode in Python 3.
  531. return torch._tensor_str._str(self, tensor_contents=tensor_contents)
  532. def backward(
  533. self, gradient=None, retain_graph=None, create_graph=False, inputs=None
  534. ):
  535. r"""Computes the gradient of current tensor wrt graph leaves.
  536. The graph is differentiated using the chain rule. If the tensor is
  537. non-scalar (i.e. its data has more than one element) and requires
  538. gradient, the function additionally requires specifying a ``gradient``.
  539. It should be a tensor of matching type and shape, that represents
  540. the gradient of the differentiated function w.r.t. ``self``.
  541. This function accumulates gradients in the leaves - you might need to zero
  542. ``.grad`` attributes or set them to ``None`` before calling it.
  543. See :ref:`Default gradient layouts<default-grad-layouts>`
  544. for details on the memory layout of accumulated gradients.
  545. .. note::
  546. If you run any forward ops, create ``gradient``, and/or call ``backward``
  547. in a user-specified CUDA stream context, see
  548. :ref:`Stream semantics of backward passes<bwd-cuda-stream-semantics>`.
  549. .. note::
  550. When ``inputs`` are provided and a given input is not a leaf,
  551. the current implementation will call its grad_fn (though it is not strictly needed to get this gradients).
  552. It is an implementation detail on which the user should not rely.
  553. See https://github.com/pytorch/pytorch/pull/60521#issuecomment-867061780 for more details.
  554. Args:
  555. gradient (Tensor, optional): The gradient of the function
  556. being differentiated w.r.t. ``self``.
  557. This argument can be omitted if ``self`` is a scalar. Defaults to ``None``.
  558. retain_graph (bool, optional): If ``False``, the graph used to compute the grads will be freed;
  559. If ``True``, it will be retained. The default is ``None``, in which case the value is inferred from ``create_graph``
  560. (i.e., the graph is retained only when higher-order derivative tracking is requested). Note that in nearly all cases
  561. setting this option to True is not needed and often can be worked around in a much more efficient way.
  562. create_graph (bool, optional): If ``True``, graph of the derivative will
  563. be constructed, allowing to compute higher order derivative
  564. products. Defaults to ``False``.
  565. inputs (Sequence[Tensor], optional): Inputs w.r.t. which the gradient will be
  566. accumulated into ``.grad``. All other tensors will be ignored. If not
  567. provided, the gradient is accumulated into all the leaf Tensors that were
  568. used to compute the :attr:`tensors`. Defaults to ``None``.
  569. """
  570. if has_torch_function_unary(self):
  571. return handle_torch_function(
  572. Tensor.backward,
  573. (self,),
  574. self,
  575. gradient=gradient,
  576. retain_graph=retain_graph,
  577. create_graph=create_graph,
  578. inputs=inputs,
  579. )
  580. torch.autograd.backward(
  581. self, gradient, retain_graph, create_graph, inputs=inputs
  582. )
  583. def index(self, positions, dims):
  584. """
  585. Index a regular tensor by binding specified positions to dims.
  586. This converts a regular tensor to a first-class tensor by binding
  587. the specified positional dimensions to Dim objects.
  588. Args:
  589. positions: Tuple of dimension positions to bind
  590. dims: Dim objects or tuple of Dim objects to bind to
  591. Returns:
  592. First-class tensor with specified dimensions bound
  593. """
  594. # TODO: make it possible to dispatch on positions/dims
  595. if has_torch_function_unary(self):
  596. return handle_torch_function(
  597. Tensor.index,
  598. (self,),
  599. self,
  600. positions,
  601. dims,
  602. )
  603. from functorch.dim import index
  604. return index(self, positions, dims)
  605. def register_hook(self, hook):
  606. r"""Registers a backward hook.
  607. The hook will be called every time a gradient with respect to the
  608. Tensor is computed. The hook should have the following signature::
  609. hook(grad) -> Tensor or None
  610. The hook should not modify its argument, but it can optionally return
  611. a new gradient which will be used in place of :attr:`grad`.
  612. This function returns a handle with a method ``handle.remove()``
  613. that removes the hook from the module.
  614. .. note::
  615. See :ref:`backward-hooks-execution` for more information on how when this hook
  616. is executed, and how its execution is ordered relative to other hooks.
  617. Example::
  618. >>> v = torch.tensor([0., 0., 0.], requires_grad=True)
  619. >>> h = v.register_hook(lambda grad: grad * 2) # double the gradient
  620. >>> v.backward(torch.tensor([1., 2., 3.]))
  621. >>> v.grad
  622. 2
  623. 4
  624. 6
  625. [torch.FloatTensor of size (3,)]
  626. >>> h.remove() # removes the hook
  627. """
  628. if has_torch_function_unary(self):
  629. return handle_torch_function(Tensor.register_hook, (self,), self, hook)
  630. if not self.requires_grad:
  631. raise RuntimeError(
  632. "cannot register a hook on a tensor that doesn't require gradient"
  633. )
  634. if self._backward_hooks is None:
  635. self._backward_hooks = OrderedDict()
  636. if self.grad_fn is not None:
  637. self.grad_fn._register_hook_dict(self)
  638. from torch.utils.hooks import RemovableHandle
  639. handle = RemovableHandle(self._backward_hooks)
  640. self._backward_hooks[handle.id] = hook
  641. return handle
  642. def register_post_accumulate_grad_hook(self, hook):
  643. r"""Registers a backward hook that runs after grad accumulation.
  644. The hook will be called after all gradients for a tensor have been accumulated,
  645. meaning that the .grad field has been updated on that tensor. The post
  646. accumulate grad hook is ONLY applicable for leaf tensors (tensors without a
  647. .grad_fn field). Registering this hook on a non-leaf tensor will error!
  648. The hook should have the following signature::
  649. hook(param: Tensor) -> None
  650. Note that, unlike other autograd hooks, this hook operates on the tensor
  651. that requires grad and not the grad itself. The hook can in-place modify
  652. and access its Tensor argument, including its .grad field.
  653. This function returns a handle with a method ``handle.remove()``
  654. that removes the hook from the module.
  655. .. note::
  656. See :ref:`backward-hooks-execution` for more information on how when this hook
  657. is executed, and how its execution is ordered relative to other hooks. Since
  658. this hook runs during the backward pass, it will run in no_grad mode (unless
  659. create_graph is True). You can use torch.enable_grad() to re-enable autograd
  660. within the hook if you need it.
  661. Example::
  662. >>> v = torch.tensor([0., 0., 0.], requires_grad=True)
  663. >>> lr = 0.01
  664. >>> # simulate a simple SGD update
  665. >>> h = v.register_post_accumulate_grad_hook(lambda p: p.add_(p.grad, alpha=-lr))
  666. >>> v.backward(torch.tensor([1., 2., 3.]))
  667. >>> v
  668. tensor([-0.0100, -0.0200, -0.0300], requires_grad=True)
  669. >>> h.remove() # removes the hook
  670. """
  671. if has_torch_function_unary(self):
  672. return handle_torch_function(
  673. Tensor.register_post_accumulate_grad_hook, (self,), self, hook
  674. )
  675. if not self.requires_grad:
  676. raise RuntimeError(
  677. "cannot register a hook on a tensor that doesn't require gradient"
  678. )
  679. if self.grad_fn is not None:
  680. raise RuntimeError(
  681. "post accumulate grad hooks cannot be registered on non-leaf tensors"
  682. )
  683. if self._post_accumulate_grad_hooks is None:
  684. self._post_accumulate_grad_hooks: dict[Any, Any] = (
  685. # pyrefly: ignore [bad-assignment]
  686. OrderedDict()
  687. )
  688. from torch.utils.hooks import RemovableHandle
  689. handle = RemovableHandle(self._post_accumulate_grad_hooks)
  690. self._post_accumulate_grad_hooks[handle.id] = hook
  691. return handle
  692. def reinforce(self, reward):
  693. def trim(str):
  694. return "\n".join([line.strip() for line in str.split("\n")])
  695. raise RuntimeError(
  696. trim(
  697. r"""reinforce() was removed.
  698. Use torch.distributions instead.
  699. See https://pytorch.org/docs/main/distributions.html
  700. Instead of:
  701. probs = policy_network(state)
  702. action = probs.multinomial()
  703. next_state, reward = env.step(action)
  704. action.reinforce(reward)
  705. action.backward()
  706. Use:
  707. probs = policy_network(state)
  708. # NOTE: categorical is equivalent to what used to be called multinomial
  709. m = torch.distributions.Categorical(probs)
  710. action = m.sample()
  711. next_state, reward = env.step(action)
  712. loss = -m.log_prob(action) * reward
  713. loss.backward()
  714. """
  715. )
  716. )
  717. detach = _C._add_docstr(
  718. _C.TensorBase.detach,
  719. r"""
  720. Returns a new Tensor, detached from the current graph.
  721. The result will never require gradient.
  722. This method also affects forward mode AD gradients and the result will never
  723. have forward mode AD gradients.
  724. .. note::
  725. Returned Tensor shares the same storage with the original one.
  726. In-place modifications on either of them will be seen, and may trigger
  727. errors in correctness checks.
  728. """,
  729. )
  730. detach_ = _C._add_docstr(
  731. _C.TensorBase.detach_,
  732. r"""
  733. Detaches the Tensor from the graph that created it, making it a leaf.
  734. Views cannot be detached in-place.
  735. This method also affects forward mode AD gradients and the result will never
  736. have forward mode AD gradients.
  737. """,
  738. )
  739. def is_shared(self):
  740. r"""Checks if tensor is in shared memory.
  741. This is always ``True`` for CUDA tensors.
  742. """
  743. if has_torch_function_unary(self):
  744. return handle_torch_function(Tensor.is_shared, (self,), self)
  745. return self._typed_storage()._is_shared()
  746. def share_memory_(self):
  747. r"""Moves the underlying storage to shared memory.
  748. This is a no-op if the underlying storage is already in shared memory
  749. and for CUDA tensors. Tensors in shared memory cannot be resized.
  750. See :meth:`torch.UntypedStorage.share_memory_` for more details.
  751. """
  752. if has_torch_function_unary(self):
  753. return handle_torch_function(Tensor.share_memory_, (self,), self)
  754. self._typed_storage()._share_memory_()
  755. return self
  756. def module_load(self, other, assign=False):
  757. r"""Defines how to transform ``other`` when loading it into ``self`` in :meth:`~nn.Module.load_state_dict`.
  758. Used when :func:`~torch.__future__.get_swap_module_params_on_conversion` is ``True``.
  759. It is expected that ``self`` is a parameter or buffer in an ``nn.Module`` and ``other`` is the
  760. value in the state dictionary with the corresponding key, this method defines
  761. how ``other`` is remapped before being swapped with ``self`` via
  762. :func:`~torch.utils.swap_tensors` in :meth:`~nn.Module.load_state_dict`.
  763. .. note::
  764. This method should always return a new object that is not ``self`` or ``other``.
  765. For example, the default implementation returns ``self.copy_(other).detach()``
  766. if ``assign`` is ``False`` or ``other.detach()`` if ``assign`` is ``True``.
  767. Args:
  768. other (Tensor): value in state dict with key corresponding to ``self``
  769. assign (bool): the assign argument passed to :meth:`nn.Module.load_state_dict`
  770. """
  771. if has_torch_function_variadic(self, other):
  772. return handle_torch_function(
  773. Tensor.module_load, (self, other), self, other, assign=assign
  774. )
  775. if assign:
  776. return other.detach()
  777. else:
  778. return self.copy_(other).detach()
  779. def __reversed__(self):
  780. r"""Reverses the tensor along dimension 0."""
  781. if has_torch_function_unary(self):
  782. return handle_torch_function(Tensor.__reversed__, (self,), self)
  783. if self.dim() == 0:
  784. return self
  785. else:
  786. return self.flip(0)
  787. def norm(
  788. self,
  789. p: float | str | None = "fro",
  790. dim=None,
  791. keepdim=False,
  792. dtype=None,
  793. ):
  794. r"""See :func:`torch.norm`"""
  795. if has_torch_function_unary(self):
  796. return handle_torch_function(
  797. Tensor.norm, (self,), self, p=p, dim=dim, keepdim=keepdim, dtype=dtype
  798. )
  799. return torch.norm(self, p, dim, keepdim, dtype=dtype)
  800. def solve(self, other):
  801. from torch._linalg_utils import solve
  802. return solve(self, other)
  803. def lstsq(self, other):
  804. from torch._linalg_utils import lstsq
  805. return lstsq(self, other)
  806. def eig(self, eigenvectors=False):
  807. from torch._linalg_utils import eig
  808. return eig(self, eigenvectors=eigenvectors)
  809. def symeig(self, eigenvectors=False):
  810. from torch._linalg_utils import _symeig
  811. return _symeig(self, eigenvectors=eigenvectors)
  812. def lu(self, pivot=True, get_infos=False):
  813. r"""See :func:`torch.lu`"""
  814. # If get_infos is True, then we don't need to check for errors and vice versa
  815. if has_torch_function_unary(self):
  816. return handle_torch_function(
  817. Tensor.lu, (self,), self, pivot=pivot, get_infos=get_infos
  818. )
  819. LU, pivots, infos = torch._lu_with_info(
  820. self, pivot=pivot, check_errors=(not get_infos)
  821. )
  822. if get_infos:
  823. return LU, pivots, infos
  824. else:
  825. return LU, pivots
  826. def stft(
  827. self,
  828. n_fft: int,
  829. hop_length: int | None = None,
  830. win_length: int | None = None,
  831. window: "Tensor | None" = None,
  832. center: bool = True,
  833. pad_mode: str = "reflect",
  834. normalized: bool = False,
  835. onesided: bool | None = None,
  836. return_complex: bool | None = None,
  837. align_to_window: bool | None = None,
  838. ):
  839. r"""See :func:`torch.stft`
  840. .. warning::
  841. This function changed signature at version 0.4.1. Calling with
  842. the previous signature may cause error or return incorrect result.
  843. """
  844. if has_torch_function_unary(self):
  845. return handle_torch_function(
  846. Tensor.stft,
  847. (self,),
  848. self,
  849. n_fft,
  850. hop_length=hop_length,
  851. win_length=win_length,
  852. window=window,
  853. center=center,
  854. pad_mode=pad_mode,
  855. normalized=normalized,
  856. onesided=onesided,
  857. return_complex=return_complex,
  858. align_to_window=align_to_window,
  859. )
  860. return torch.stft(
  861. self,
  862. n_fft,
  863. hop_length,
  864. win_length,
  865. window,
  866. center,
  867. pad_mode,
  868. normalized,
  869. onesided,
  870. return_complex=return_complex,
  871. align_to_window=align_to_window,
  872. )
  873. def istft(
  874. self,
  875. n_fft: int,
  876. hop_length: int | None = None,
  877. win_length: int | None = None,
  878. window: "Tensor | None" = None,
  879. center: bool = True,
  880. normalized: bool = False,
  881. onesided: bool | None = None,
  882. length: int | None = None,
  883. return_complex: bool = False,
  884. ):
  885. r"""See :func:`torch.istft`"""
  886. if has_torch_function_unary(self):
  887. return handle_torch_function(
  888. Tensor.istft,
  889. (self,),
  890. self,
  891. n_fft,
  892. hop_length=hop_length,
  893. win_length=win_length,
  894. window=window,
  895. center=center,
  896. normalized=normalized,
  897. onesided=onesided,
  898. length=length,
  899. return_complex=return_complex,
  900. )
  901. return torch.istft(
  902. self,
  903. n_fft,
  904. hop_length,
  905. win_length,
  906. window,
  907. center,
  908. normalized,
  909. onesided,
  910. length,
  911. return_complex=return_complex,
  912. )
  913. def resize(self, *sizes):
  914. if has_torch_function_unary(self):
  915. return handle_torch_function(Tensor.resize, (self,), self, *sizes)
  916. warnings.warn("non-inplace resize is deprecated", stacklevel=2)
  917. from torch.autograd._functions import Resize
  918. return Resize.apply(self, sizes)
  919. def resize_as(self, tensor):
  920. if has_torch_function_variadic(self, tensor):
  921. return handle_torch_function(Tensor.resize_as, (self, tensor), self, tensor)
  922. warnings.warn("non-inplace resize_as is deprecated", stacklevel=2)
  923. from torch.autograd._functions import Resize
  924. return Resize.apply(self, tensor.size())
  925. def split(self, split_size, dim=0):
  926. r"""See :func:`torch.split`"""
  927. if has_torch_function_unary(self):
  928. return handle_torch_function(
  929. Tensor.split, (self,), self, split_size, dim=dim
  930. )
  931. if isinstance(split_size, Tensor):
  932. try:
  933. split_size = int(split_size)
  934. except ValueError:
  935. pass
  936. if isinstance(split_size, (int, torch.SymInt)):
  937. return torch._VF.split(self, split_size, dim) # type: ignore[attr-defined]
  938. else:
  939. return torch._VF.split_with_sizes(
  940. self,
  941. split_size,
  942. dim,
  943. )
  944. def unique(self, sorted=True, return_inverse=False, return_counts=False, dim=None):
  945. r"""Returns the unique elements of the input tensor.
  946. See :func:`torch.unique`
  947. """
  948. if has_torch_function_unary(self):
  949. return handle_torch_function(
  950. Tensor.unique,
  951. (self,),
  952. self,
  953. sorted=sorted,
  954. return_inverse=return_inverse,
  955. return_counts=return_counts,
  956. dim=dim,
  957. )
  958. return torch.unique(
  959. self,
  960. sorted=sorted,
  961. return_inverse=return_inverse,
  962. return_counts=return_counts,
  963. dim=dim,
  964. )
  965. def unique_consecutive(self, return_inverse=False, return_counts=False, dim=None):
  966. r"""Eliminates all but the first element from every consecutive group of equivalent elements.
  967. See :func:`torch.unique_consecutive`
  968. """
  969. if has_torch_function_unary(self):
  970. return handle_torch_function(
  971. Tensor.unique_consecutive,
  972. (self,),
  973. self,
  974. return_inverse=return_inverse,
  975. return_counts=return_counts,
  976. dim=dim,
  977. )
  978. return torch.unique_consecutive(
  979. self, return_inverse=return_inverse, return_counts=return_counts, dim=dim
  980. )
  981. @_handle_torch_function_and_wrap_type_error_to_not_implemented
  982. def __rsub__(self, other: Union["Tensor", int, float, bool, complex]) -> "Tensor":
  983. return _C._VariableFunctions.rsub(self, other)
  984. @_handle_torch_function_and_wrap_type_error_to_not_implemented
  985. def __rdiv__(self, other: Union["Tensor", int, float, bool, complex]) -> "Tensor":
  986. return self.reciprocal() * other
  987. __rtruediv__ = __rdiv__
  988. __itruediv__ = _C.TensorBase.__idiv__
  989. # pyrefly: ignore [bad-override]
  990. __pow__ = cast(
  991. Callable[
  992. ["torch._C.TensorBase", Union["Tensor", int, float, bool, complex]],
  993. "Tensor",
  994. ],
  995. _handle_torch_function_and_wrap_type_error_to_not_implemented(
  996. _C.TensorBase.pow
  997. ),
  998. )
  999. __ipow__ = _handle_torch_function_and_wrap_type_error_to_not_implemented(
  1000. _C.TensorBase.pow_
  1001. )
  1002. @_handle_torch_function_and_wrap_type_error_to_not_implemented
  1003. def __rmod__(self, other: Union["Tensor", int, float, bool, complex]) -> "Tensor":
  1004. return torch.remainder(other, self)
  1005. def __format__(self, format_spec):
  1006. if has_torch_function_unary(self):
  1007. return handle_torch_function(Tensor.__format__, (self,), self, format_spec)
  1008. if self.dim() == 0 and not self.is_meta and type(self) is Tensor:
  1009. # Use detach() here to avoid the warning when converting a scalar Tensor that
  1010. # requires gradients to a python number. It is ok for formatting.
  1011. return self.detach().item().__format__(format_spec)
  1012. return object.__format__(self, format_spec)
  1013. @_handle_torch_function_and_wrap_type_error_to_not_implemented
  1014. def __rpow__(self, other: Union["Tensor", int, float, bool, complex]) -> "Tensor":
  1015. return torch.pow(other, self)
  1016. @_handle_torch_function_and_wrap_type_error_to_not_implemented
  1017. def __floordiv__(self, other: Union["Tensor", int, float, bool]) -> "Tensor": # type: ignore[override]
  1018. # TODO(rec): the superclass says it accepts complex here,
  1019. # but torch.floor_divide says it doesn't.
  1020. return torch.floor_divide(self, other)
  1021. @_handle_torch_function_and_wrap_type_error_to_not_implemented
  1022. def __rfloordiv__(self, other: Union["Tensor", int, float, bool]) -> "Tensor": # type: ignore[override]
  1023. return torch.floor_divide(other, self)
  1024. @_handle_torch_function_and_wrap_type_error_to_not_implemented
  1025. def __rlshift__(
  1026. self, other: Union["Tensor", int, float, bool, complex]
  1027. ) -> "Tensor":
  1028. return torch.bitwise_left_shift(other, self)
  1029. @_handle_torch_function_and_wrap_type_error_to_not_implemented
  1030. def __rrshift__(
  1031. self, other: Union["Tensor", int, float, bool, complex]
  1032. ) -> "Tensor":
  1033. return torch.bitwise_right_shift(other, self)
  1034. @_handle_torch_function_and_wrap_type_error_to_not_implemented
  1035. def __rmatmul__(self, other: "Tensor") -> "Tensor":
  1036. return torch.matmul(other, self)
  1037. __pos__ = _C.TensorBase.positive
  1038. __neg__ = _C.TensorBase.neg
  1039. __abs__ = _C.TensorBase.abs
  1040. def __len__(self):
  1041. if has_torch_function_unary(self):
  1042. return handle_torch_function(Tensor.__len__, (self,), self)
  1043. if self.dim() == 0:
  1044. raise TypeError("len() of a 0-d tensor")
  1045. if torch._C._get_tracing_state():
  1046. warnings.warn(
  1047. "Using len to get tensor shape might cause the trace to be incorrect. "
  1048. "Recommended usage would be tensor.shape[0]. "
  1049. "Passing a tensor of different shape might lead to errors or silently give "
  1050. "incorrect results.",
  1051. category=torch.jit.TracerWarning,
  1052. stacklevel=2,
  1053. )
  1054. return self.shape[0]
  1055. def __iter__(self):
  1056. # NB: we use 'imap' and not 'map' here, so that in Python 2 we get a
  1057. # generator and don't eagerly perform all the indexes. This could
  1058. # save us work, and also helps keep trace ordering deterministic
  1059. # (e.g., if you zip(*hiddens), the eager map will force all the
  1060. # indexes of hiddens[0] before hiddens[1], while the generator
  1061. # map will interleave them.)
  1062. # NB: We have intentionally skipped __torch_function__ dispatch here.
  1063. # See gh-54457
  1064. if self.dim() == 0:
  1065. raise TypeError("iteration over a 0-d tensor")
  1066. if torch._C._get_tracing_state():
  1067. warnings.warn(
  1068. "Iterating over a tensor might cause the trace to be incorrect. "
  1069. "Passing a tensor of different shape won't change the number of "
  1070. "iterations executed (and might lead to errors or silently give "
  1071. "incorrect results).",
  1072. category=torch.jit.TracerWarning,
  1073. stacklevel=2,
  1074. )
  1075. return iter(self.unbind(0))
  1076. def __hash__(self):
  1077. # Do NOT handle __torch_function__ here as user's default
  1078. # implementation that handle most functions will most likely do it wrong.
  1079. # It can be easily overridden by defining this method on the user
  1080. # subclass if needed.
  1081. return id(self)
  1082. def __dir__(self):
  1083. if has_torch_function_unary(self):
  1084. return handle_torch_function(Tensor.__dir__, (self,), self)
  1085. tensor_methods = dir(self.__class__)
  1086. tensor_methods.remove("volatile") # deprecated
  1087. attrs = list(self.__dict__.keys())
  1088. keys = tensor_methods + attrs
  1089. # property only available dense, cuda tensors
  1090. if (not self.is_cuda) or self.is_sparse:
  1091. keys.remove("__cuda_array_interface__")
  1092. return sorted(keys)
  1093. # Numpy array interface, to support `numpy.asarray(tensor) -> ndarray`
  1094. __array_priority__ = 1000 # prefer Tensor ops over numpy ones
  1095. def __array__(self, dtype=None):
  1096. if has_torch_function_unary(self):
  1097. return handle_torch_function(Tensor.__array__, (self,), self, dtype=dtype)
  1098. if dtype is None:
  1099. return self.numpy()
  1100. else:
  1101. return self.numpy().astype(dtype, copy=False)
  1102. # Wrap Numpy array again in a suitable tensor when done, to support e.g.
  1103. # `numpy.sin(tensor) -> tensor` or `numpy.greater(tensor, 0) -> ByteTensor`
  1104. def __array_wrap__(self, array):
  1105. if has_torch_function_unary(self):
  1106. return handle_torch_function(
  1107. Tensor.__array_wrap__, (self,), self, array=array
  1108. )
  1109. if array.dtype == bool:
  1110. # Workaround, torch has no built-in bool tensor
  1111. array = array.astype("uint8")
  1112. return torch.from_numpy(array)
  1113. def __contains__(self, element: Any, /) -> bool:
  1114. r"""Check if `element` is present in tensor
  1115. Args:
  1116. element (Tensor or scalar): element to be checked
  1117. for presence in current tensor"
  1118. """
  1119. if has_torch_function_unary(self):
  1120. return handle_torch_function(Tensor.__contains__, (self,), self, element)
  1121. if isinstance(
  1122. element, (torch.Tensor, Number, torch.SymInt, torch.SymFloat, torch.SymBool)
  1123. ):
  1124. # type hint doesn't understand the __contains__ result array
  1125. return bool((element == self).any().item()) # type: ignore[union-attr]
  1126. raise RuntimeError(
  1127. f"Tensor.__contains__ only supports Tensor or scalar, but you passed in a {type(element)}."
  1128. )
  1129. @property
  1130. def __cuda_array_interface__(self):
  1131. """Array view description for cuda tensors.
  1132. See:
  1133. https://numba.pydata.org/numba-doc/dev/cuda/cuda_array_interface.html
  1134. """
  1135. if has_torch_function_unary(self):
  1136. # TODO mypy doesn't support @property, see: https://github.com/python/mypy/issues/6185
  1137. return handle_torch_function(
  1138. Tensor.__cuda_array_interface__.__get__, # type: ignore[attr-defined]
  1139. (self,),
  1140. self,
  1141. )
  1142. # raise AttributeError for unsupported tensors, so that
  1143. # hasattr(cpu_tensor, "__cuda_array_interface__") is False.
  1144. if not self.is_cuda:
  1145. raise AttributeError(
  1146. f"Can't get __cuda_array_interface__ on non-CUDA tensor type: {self.type()} "
  1147. "If CUDA data is required use tensor.cuda() to copy tensor to device memory."
  1148. )
  1149. if self.is_sparse:
  1150. raise AttributeError(
  1151. f"Can't get __cuda_array_interface__ on sparse type: {self.type()} "
  1152. "Use Tensor.to_dense() to convert to a dense tensor first."
  1153. )
  1154. # RuntimeError, matching tensor.__array__() behavior.
  1155. if self.requires_grad:
  1156. raise RuntimeError(
  1157. "Can't get __cuda_array_interface__ on Variable that requires grad. "
  1158. "If gradients aren't required, use var.detach() to get Variable that doesn't require grad."
  1159. )
  1160. typestr = _dtype_to_typestr(self.dtype)
  1161. itemsize = self.element_size()
  1162. shape = tuple(self.shape)
  1163. if self.is_contiguous():
  1164. # __cuda_array_interface__ v2 requires the strides to be omitted
  1165. # (either not set or set to None) for C-contiguous arrays.
  1166. strides = None
  1167. else:
  1168. strides = tuple(s * itemsize for s in self.stride())
  1169. data_ptr = self.data_ptr() if self.numel() > 0 else 0
  1170. data = (data_ptr, False) # read-only is false
  1171. return dict(typestr=typestr, shape=shape, strides=strides, data=data, version=2)
  1172. def storage_type(self):
  1173. r"""storage_type() -> type
  1174. Returns the type of the underlying storage.
  1175. """
  1176. if has_torch_function_unary(self):
  1177. return handle_torch_function(Tensor.storage_type, (self,), self)
  1178. torch.storage._warn_typed_storage_removal()
  1179. return self._typed_storage()._get_legacy_storage_class()
  1180. def refine_names(self, *names): # pyrefly: ignore # bad-override
  1181. r"""Refines the dimension names of :attr:`self` according to :attr:`names`.
  1182. Refining is a special case of renaming that "lifts" unnamed dimensions.
  1183. A ``None`` dim can be refined to have any name; a named dim can only be
  1184. refined to have the same name.
  1185. Because named tensors can coexist with unnamed tensors, refining names
  1186. gives a nice way to write named-tensor-aware code that works with both
  1187. named and unnamed tensors.
  1188. :attr:`names` may contain up to one Ellipsis (``...``).
  1189. The Ellipsis is expanded greedily; it is expanded in-place to fill
  1190. :attr:`names` to the same length as ``self.dim()`` using names from the
  1191. corresponding indices of ``self.names``.
  1192. Python 2 does not support Ellipsis but one may use a string literal
  1193. instead (``'...'``).
  1194. Args:
  1195. names (iterable of str): The desired names of the output tensor. May
  1196. contain up to one Ellipsis.
  1197. Examples::
  1198. >>> imgs = torch.randn(32, 3, 128, 128)
  1199. >>> named_imgs = imgs.refine_names('N', 'C', 'H', 'W')
  1200. >>> named_imgs.names
  1201. ('N', 'C', 'H', 'W')
  1202. >>> tensor = torch.randn(2, 3, 5, 7, 11)
  1203. >>> tensor = tensor.refine_names('A', ..., 'B', 'C')
  1204. >>> tensor.names
  1205. ('A', None, None, 'B', 'C')
  1206. .. warning::
  1207. The named tensor API is experimental and subject to change.
  1208. """
  1209. if has_torch_function_unary(self):
  1210. return handle_torch_function(Tensor.refine_names, (self,), self, *names)
  1211. names = resolve_ellipsis(names, self.names, "refine_names")
  1212. return super().refine_names(names)
  1213. def align_to(self, *names): # pyrefly: ignore # bad-override
  1214. r"""Permutes the dimensions of the :attr:`self` tensor to match the order
  1215. specified in :attr:`names`, adding size-one dims for any new names.
  1216. All of the dims of :attr:`self` must be named in order to use this method.
  1217. The resulting tensor is a view on the original tensor.
  1218. All dimension names of :attr:`self` must be present in :attr:`names`.
  1219. :attr:`names` may contain additional names that are not in ``self.names``;
  1220. the output tensor has a size-one dimension for each of those new names.
  1221. :attr:`names` may contain up to one Ellipsis (``...``).
  1222. The Ellipsis is expanded to be equal to all dimension names of :attr:`self`
  1223. that are not mentioned in :attr:`names`, in the order that they appear
  1224. in :attr:`self`.
  1225. Python 2 does not support Ellipsis but one may use a string literal
  1226. instead (``'...'``).
  1227. Args:
  1228. names (iterable of str): The desired dimension ordering of the
  1229. output tensor. May contain up to one Ellipsis that is expanded
  1230. to all unmentioned dim names of :attr:`self`.
  1231. Examples::
  1232. >>> tensor = torch.randn(2, 2, 2, 2, 2, 2)
  1233. >>> named_tensor = tensor.refine_names('A', 'B', 'C', 'D', 'E', 'F')
  1234. # Move the F and E dims to the front while keeping the rest in order
  1235. >>> named_tensor.align_to('F', 'E', ...)
  1236. .. warning::
  1237. The named tensor API is experimental and subject to change.
  1238. """
  1239. if has_torch_function_unary(self):
  1240. return handle_torch_function(Tensor.align_to, (self,), self, *names)
  1241. ellipsis_idx = single_ellipsis_index(names, "align_to")
  1242. if ellipsis_idx is None:
  1243. return super().align_to(names)
  1244. return super().align_to(
  1245. [name for name in names if not is_ellipsis(name)], ellipsis_idx
  1246. )
  1247. def unflatten(self, dim, sizes): # type: ignore[override]
  1248. r"""
  1249. unflatten(dim, sizes) -> Tensor
  1250. See :func:`torch.unflatten`.
  1251. """
  1252. if has_torch_function_unary(self):
  1253. return handle_torch_function(Tensor.unflatten, (self,), self, dim, sizes)
  1254. if not sizes:
  1255. raise RuntimeError("unflatten: sizes must be non-empty")
  1256. names = None
  1257. if isinstance(sizes, OrderedDict) or (
  1258. isinstance(sizes, (tuple, list)) and isinstance(sizes[0], (tuple, list))
  1259. ):
  1260. names, sizes = unzip_namedshape(sizes)
  1261. return super().unflatten(dim, sizes, names)
  1262. else:
  1263. return super().unflatten(dim, sizes)
  1264. def rename_(self, *names, **rename_map):
  1265. """In-place version of :meth:`~Tensor.rename`."""
  1266. if has_torch_function_unary(self):
  1267. return handle_torch_function(
  1268. Tensor.rename_, (self,), self, *names, **rename_map
  1269. )
  1270. # Note [rename_ / rename API]
  1271. # The Python API for these is different from the C++ API. In Python:
  1272. # 1) tensor.rename(*names) takes a vararglist of names
  1273. # 2) tensor.rename(**rename_map) takes a map of names to rename.
  1274. # C++ is static, making it difficult to implement similar behavior.
  1275. return update_names(self, names, rename_map, inplace=True)
  1276. def rename(self, *names, **rename_map):
  1277. """Renames dimension names of :attr:`self`.
  1278. There are two main usages:
  1279. ``self.rename(**rename_map)`` returns a view on tensor that has dims
  1280. renamed as specified in the mapping :attr:`rename_map`.
  1281. ``self.rename(*names)`` returns a view on tensor, renaming all
  1282. dimensions positionally using :attr:`names`.
  1283. Use ``self.rename(None)`` to drop names on a tensor.
  1284. One cannot specify both positional args :attr:`names` and keyword args
  1285. :attr:`rename_map`.
  1286. Examples::
  1287. >>> imgs = torch.rand(2, 3, 5, 7, names=('N', 'C', 'H', 'W'))
  1288. >>> renamed_imgs = imgs.rename(N='batch', C='channels')
  1289. >>> renamed_imgs.names
  1290. ('batch', 'channels', 'H', 'W')
  1291. >>> renamed_imgs = imgs.rename(None)
  1292. >>> renamed_imgs.names
  1293. (None, None, None, None)
  1294. >>> renamed_imgs = imgs.rename('batch', 'channel', 'height', 'width')
  1295. >>> renamed_imgs.names
  1296. ('batch', 'channel', 'height', 'width')
  1297. .. warning::
  1298. The named tensor API is experimental and subject to change.
  1299. """
  1300. if has_torch_function_unary(self):
  1301. return handle_torch_function(
  1302. Tensor.rename, (self,), self, *names, **rename_map
  1303. )
  1304. # See Note [rename_ / rename API]
  1305. return update_names(self, names, rename_map, inplace=False)
  1306. def to_sparse_coo(self):
  1307. """Convert a tensor to :ref:`coordinate format <sparse-coo-docs>`.
  1308. Examples::
  1309. >>> dense = torch.randn(5, 5)
  1310. >>> sparse = dense.to_sparse_coo()
  1311. >>> sparse._nnz()
  1312. 25
  1313. """
  1314. return self.to_sparse()
  1315. def dim_order(self, *, ambiguity_check: bool | list[torch.memory_format] = False):
  1316. """
  1317. dim_order(ambiguity_check=False) -> tuple
  1318. Returns the uniquely determined tuple of int describing the dim order or
  1319. physical layout of :attr:`self`.
  1320. The dim order represents how dimensions are laid out in memory of dense tensors,
  1321. starting from the outermost to the innermost dimension.
  1322. Note that the dim order may not always be uniquely determined.
  1323. If `ambiguity_check` is True, this function raises a RuntimeError when the dim order cannot be uniquely determined;
  1324. If `ambiguity_check` is a list of memory formats, this function raises a RuntimeError when tensor can not be interpreted
  1325. into exactly one of the given memory formats, or it cannot be uniquely determined.
  1326. If `ambiguity_check` is False, it will return one of legal dim order(s) without checking its uniqueness.
  1327. Otherwise, it will raise TypeError.
  1328. Args:
  1329. ambiguity_check (bool or List[torch.memory_format]): The check method for ambiguity of dim order.
  1330. Examples::
  1331. >>> torch.empty((2, 3, 5, 7)).dim_order()
  1332. (0, 1, 2, 3)
  1333. >>> torch.empty((2, 3, 5, 7)).transpose(1, 2).dim_order()
  1334. (0, 2, 1, 3)
  1335. >>> torch.empty((2, 3, 5, 7), memory_format=torch.channels_last).dim_order()
  1336. (0, 2, 3, 1)
  1337. >>> torch.empty((1, 2, 3, 4)).dim_order()
  1338. (0, 1, 2, 3)
  1339. >>> try:
  1340. ... torch.empty((1, 2, 3, 4)).dim_order(ambiguity_check=True)
  1341. ... except RuntimeError as e:
  1342. ... print(e)
  1343. The tensor does not have unique dim order, or cannot map to exact one of the given memory formats.
  1344. >>> torch.empty((1, 2, 3, 4)).dim_order(
  1345. ... ambiguity_check=[torch.contiguous_format, torch.channels_last]
  1346. ... ) # It can be mapped to contiguous format
  1347. (0, 1, 2, 3)
  1348. >>> try:
  1349. ... torch.empty((1, 2, 3, 4)).dim_order(ambiguity_check="ILLEGAL") # type: ignore[arg-type]
  1350. ... except TypeError as e:
  1351. ... print(e)
  1352. The ambiguity_check argument must be a bool or a list of memory formats.
  1353. .. warning::
  1354. The dim_order tensor API is experimental and subject to change.
  1355. """
  1356. if has_torch_function_unary(self):
  1357. return handle_torch_function(Tensor.dim_order, (self,), self)
  1358. if self.is_sparse:
  1359. raise AttributeError(
  1360. f"Can't get dim order on sparse type: {self.type()} "
  1361. "Use Tensor.to_dense() to convert to a dense tensor first."
  1362. )
  1363. # Sanity check ambiguity_check data types
  1364. if not isinstance(ambiguity_check, bool):
  1365. if not isinstance(ambiguity_check, list):
  1366. raise TypeError(
  1367. "The ambiguity_check argument must be a bool or a list of memory formats."
  1368. )
  1369. for memory_format in ambiguity_check:
  1370. if not isinstance(memory_format, torch.memory_format):
  1371. raise TypeError(
  1372. "The ambiguity_check argument must be a bool or a list of memory formats."
  1373. )
  1374. def invalid_unique_memory_format(tensor, valid_memory_formats):
  1375. """
  1376. Returns True if the tensor cannot be uniquely mapped to any of the given memory formats, False otherwise.
  1377. """
  1378. n_legality = 0
  1379. for memory_format in valid_memory_formats:
  1380. if tensor.is_contiguous(memory_format=memory_format):
  1381. n_legality += 1
  1382. return n_legality != 1
  1383. def has_multiple_dim_order(tensor):
  1384. """
  1385. Returns True if there're multiple legal dim orders for given tensor, False otherwise.
  1386. The tensor is considered to have multiple legal dim orders if either of the following conditions is met:
  1387. * Singleton Dimensions: There's at least one singleteon dimension in the tensor.
  1388. Since their size is 1, they don't affect the memory offset (stride * index
  1389. is zero because index is always zero). Therefore, they can be placed anywhere
  1390. in the dimension order without changing how data is accessed.
  1391. * Same strides: Strides reflect how the tensor is stored in memory.
  1392. If any two dimensions have the same stride, swapping these dimensions won't
  1393. change how data is accessed, leading to multiple correct dimension orders.
  1394. """
  1395. from torch.fx.experimental.symbolic_shapes import guard_or_false
  1396. sizes = tensor.size()
  1397. strides = tensor.stride()
  1398. # Check if there are any duplicate strides
  1399. has_duplicate_strides = any(
  1400. guard_or_false(earlier == later)
  1401. for earlier, later in itertools.pairwise(strides)
  1402. )
  1403. # Check if there are any singleton dimensions
  1404. has_singleton_dims = any(guard_or_false(size == 1) for size in sizes)
  1405. return has_duplicate_strides or has_singleton_dims
  1406. valid_memory_formats = (
  1407. ambiguity_check if isinstance(ambiguity_check, list) else []
  1408. )
  1409. check_multiple_dim_order = (
  1410. ambiguity_check if isinstance(ambiguity_check, bool) else True
  1411. )
  1412. if (
  1413. check_multiple_dim_order and has_multiple_dim_order(self)
  1414. ) and invalid_unique_memory_format(self, valid_memory_formats):
  1415. raise RuntimeError(
  1416. "The tensor does not have unique dim order, or cannot map to exact one of the given memory formats."
  1417. )
  1418. import torch._prims_common as utils
  1419. out_perm, raise_ambiguity = (
  1420. utils.compute_elementwise_output_logical_to_physical_perm(
  1421. self, ambiguity_check=ambiguity_check
  1422. )
  1423. )
  1424. if raise_ambiguity:
  1425. raise RuntimeError("The tensor does not have unique dim order.")
  1426. return tuple(out_perm)
  1427. def _update_names(self, names, inplace):
  1428. if has_torch_function_unary(self):
  1429. return handle_torch_function(
  1430. Tensor._update_names, (self,), self, names, inplace
  1431. )
  1432. # See Note [rename_ / rename API]
  1433. if inplace:
  1434. return super().rename_(names)
  1435. else:
  1436. return super().rename(names)
  1437. @classmethod
  1438. def __torch_function__(cls, func, types, args=(), kwargs=None):
  1439. """
  1440. This __torch_function__ implementation wraps subclasses such that
  1441. methods called on subclasses return a subclass instance instead of
  1442. a ``torch.Tensor`` instance.
  1443. One corollary to this is that you need coverage for torch.Tensor
  1444. methods if implementing __torch_function__ for subclasses.
  1445. We recommend always calling ``super().__torch_function__`` as the base
  1446. case when doing the above.
  1447. While not mandatory, we recommend making `__torch_function__` a classmethod.
  1448. """
  1449. if kwargs is None:
  1450. kwargs = {}
  1451. if not all(issubclass(cls, t) for t in types):
  1452. return NotImplemented
  1453. with _C.DisableTorchFunctionSubclass():
  1454. ret = func(*args, **kwargs)
  1455. if func in get_default_nowrap_functions():
  1456. return ret
  1457. else:
  1458. return _convert(ret, cls)
  1459. __torch_dispatch__ = _C._disabled_torch_dispatch_impl
  1460. def __dlpack__(
  1461. self,
  1462. *,
  1463. stream: Any | None = -1,
  1464. max_version: tuple[int, int] | None = None,
  1465. dl_device: tuple[enum.IntEnum, int] | None = None,
  1466. copy: bool | None = None,
  1467. ):
  1468. """
  1469. Creates a DLpack `capsule https://data-apis.org/array-api/latest/design_topics/data_interchange.html#data-interchange`_
  1470. of the current tensor to be exported to other libraries.
  1471. This function will be called from the `from_dlpack` method
  1472. of the library that will consume the capsule. `from_dlpack` passes the current
  1473. stream to this method as part of the specification.
  1474. Args:
  1475. stream (integer or None): An optional Python integer representing a
  1476. pointer to a CUDA stream. The current stream is synchronized with
  1477. this stream before the capsule is created, and since the capsule
  1478. shares its storage with the tensor this make it safe to access from
  1479. both streams. If -1 is passed then no synchronization is performed.
  1480. If 1 (on CUDA) or 0 (on ROCM) then the default stream is used for
  1481. synchronization. This API intentionally slightly deviates from the DLPack
  1482. guidance: the default stream is -1 (stream-preserving; no cross-stream sync),
  1483. because many from_dlpack implementations intend stream preservation.
  1484. For non-CUDA devices, -1 is treated the same as None.
  1485. max_version (tuple[int, int] or None): An optional Python tuple with
  1486. 2 integers, representing the maximum version the caller supports. If
  1487. None (default), PyTorch will fallback to DLPack 0.8.
  1488. dl_device (tuple[DLDeviceType, int] or None): An optional tuple specifying
  1489. in which device the exported DLPack capsule should be on. If None (default),
  1490. the exported DLPack capsule will be on the same device as ``self``.
  1491. copy (bool or None): An optional boolean indicating whether or not to copy
  1492. ``self``. If None, PyTorch will copy only if necessary.
  1493. """
  1494. if has_torch_function_unary(self):
  1495. args = (self,)
  1496. kwargs = {
  1497. "stream": stream,
  1498. "max_version": max_version,
  1499. "dl_device": dl_device,
  1500. "copy": copy,
  1501. }
  1502. return handle_torch_function(Tensor.__dlpack__, (self,), *args, **kwargs)
  1503. # DLPack capsules can't capture all of PyTorch's semantics,
  1504. # so we prohibit exporting tensors that would lose their properties like
  1505. # requires_grad and having the conjugate bit set.
  1506. if self.requires_grad:
  1507. raise BufferError(
  1508. "Can't export tensors that require gradient, use tensor.detach()"
  1509. )
  1510. if self.is_conj():
  1511. raise BufferError("Can't export tensors with the conjugate bit set")
  1512. if self.layout != torch.strided:
  1513. raise BufferError(
  1514. "Can't export tensors with layout other than torch.strided"
  1515. )
  1516. if (
  1517. self.device.type == "cuda"
  1518. and self.device.index != torch.cuda.current_device()
  1519. ):
  1520. raise BufferError(
  1521. "Can't export tensors on a different CUDA device index. "
  1522. f"Expected: {self.device.index}. "
  1523. f"Current device: {torch.cuda.current_device()}."
  1524. )
  1525. if stream is not None and type(stream) is not int:
  1526. # Stream pointers in CUDA/ROCm are uniquely numbered and can
  1527. # be retrieved from their integer value.
  1528. raise TypeError("stream must be ``int`` or ``none``")
  1529. elif self.device.type == "cuda" and stream != -1:
  1530. # NB: This logic handles the special case values for default
  1531. # streams and must be kept in sync with from_dlpack in
  1532. # torch/utils/dlpack.py
  1533. is_rocm = torch.version.hip is not None
  1534. is_cuda = not is_rocm
  1535. if stream is None or (is_rocm and stream == 0) or (is_cuda and stream == 1):
  1536. stream = torch.cuda.default_stream()
  1537. else:
  1538. if is_cuda and stream == 2:
  1539. raise BufferError("per-thread default stream is not supported.")
  1540. device_str = "CUDA" if is_cuda else "ROCm"
  1541. if not (
  1542. (is_cuda and stream != 0) or (is_rocm and stream not in (1, 2))
  1543. ):
  1544. raise AssertionError(
  1545. f"unsupported stream on {device_str}: {stream}."
  1546. )
  1547. stream = torch.cuda.ExternalStream(stream)
  1548. # Only synchronize on different streams
  1549. current_stream = torch.cuda.current_stream()
  1550. if stream != current_stream:
  1551. event = torch.cuda.Event()
  1552. event.record(current_stream)
  1553. stream.wait_event(event)
  1554. elif self.device.type == "cpu":
  1555. if stream is not None and stream != -1:
  1556. raise AssertionError("stream should be None on cpu.")
  1557. if self.device.type == "xla":
  1558. import torch_xla
  1559. import torch_xla.utils.dlpack as xla_dlpack
  1560. if (
  1561. len(torch_xla.real_devices()) <= 0
  1562. or "cuda" not in torch_xla.real_devices()[0].lower()
  1563. ):
  1564. raise RuntimeError(
  1565. "Can't export to dlpack an XLA tensor that is not on CUDA."
  1566. )
  1567. # Does not support DLPack 1.0, yet.
  1568. return xla_dlpack.to_dlpack(self)
  1569. if max_version is None or max_version[0] < 1:
  1570. # Fallback to the old, unversioned variant.
  1571. return _C._to_dlpack(self, dl_device=dl_device, copy=copy)
  1572. return _C._to_dlpack_versioned(self, dl_device=dl_device, copy=copy)
  1573. def __dlpack_device__(self) -> tuple[enum.IntEnum, int]:
  1574. if has_torch_function_unary(self):
  1575. return handle_torch_function(Tensor.__dlpack_device__, (self,), self)
  1576. from torch.utils.dlpack import DLDeviceType
  1577. device = self.device
  1578. idx = device.index if device.index is not None else 0
  1579. torch_device_type = device.type
  1580. if torch_device_type == "cuda" and torch.version.hip is not None:
  1581. device_type = DLDeviceType.kDLROCM
  1582. elif torch_device_type == "cpu" and self.is_pinned():
  1583. device_type = DLDeviceType.kDLCUDAHost
  1584. elif torch_device_type == "cuda":
  1585. device_type = DLDeviceType.kDLCUDA
  1586. elif torch_device_type == "cpu":
  1587. device_type = DLDeviceType.kDLCPU
  1588. elif torch_device_type == "xpu":
  1589. device_type = DLDeviceType.kDLOneAPI
  1590. elif self.device.type == "privateuse1":
  1591. device_type = DLDeviceType.kDLExtDev
  1592. elif torch_device_type == "xla":
  1593. import torch_xla
  1594. if (
  1595. len(torch_xla.real_devices()) <= 0
  1596. or "cuda" not in torch_xla.real_devices()[0].lower()
  1597. ):
  1598. raise ValueError(f"Unknown device type {torch_device_type} for Dlpack")
  1599. device_type = DLDeviceType.kDLCUDA
  1600. elif torch_device_type == "mps":
  1601. device_type = DLDeviceType.kDLMetal
  1602. else:
  1603. raise ValueError(f"Unknown device type {torch_device_type} for Dlpack")
  1604. return (device_type, idx)
  1605. __module__ = "torch"
  1606. def _convert(ret, cls):
  1607. if cls is Tensor:
  1608. return ret
  1609. if isinstance(ret, Tensor) and not isinstance(ret, cls):
  1610. ret = ret.as_subclass(cls)
  1611. if isinstance(ret, (tuple, list)):
  1612. # Also handles things like namedtuples
  1613. ret = type(ret)(_convert(r, cls) for r in ret)
  1614. return ret