__init__.py 87 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263
  1. # mypy: allow-untyped-defs
  2. r"""
  3. The torch package contains data structures for multi-dimensional
  4. tensors and defines mathematical operations over these tensors.
  5. Additionally, it provides many utilities for efficient serialization of
  6. Tensors and arbitrary types, and other useful utilities.
  7. It has a CUDA counterpart, that enables you to run your tensor computations
  8. on an NVIDIA GPU with compute capability >= 3.0.
  9. """
  10. import math
  11. import os
  12. import sys
  13. import platform
  14. import textwrap
  15. import ctypes
  16. import inspect
  17. import threading
  18. import pdb
  19. import importlib
  20. import importlib.util
  21. # multipy/deploy is setting this import before importing torch, this is the most
  22. # reliable way we have to detect if we're running within deploy.
  23. # https://github.com/pytorch/multipy/blob/d60f34ad38c371e441fe7ffdb77a3c3dda5a5d19/multipy/runtime/interpreter/interpreter_impl.cpp#L134-L137
  24. def _running_with_deploy():
  25. return sys.modules.get("torch._meta_registrations", None) is object
  26. from ._utils import _import_dotted_name, classproperty
  27. from ._utils import _functionalize_sync as _sync
  28. from ._utils_internal import get_file_path, prepare_multiprocessing_environment, \
  29. USE_RTLD_GLOBAL_WITH_LIBTORCH, USE_GLOBAL_DEPS
  30. # TODO(torch_deploy) figure out how to freeze version.py in fbcode build
  31. if _running_with_deploy():
  32. __version__ = "torch-deploy-1.8"
  33. else:
  34. from .torch_version import __version__ as __version__
  35. from typing import Any, Callable, Dict, Optional, Set, Tuple, Type, TYPE_CHECKING, Union, List
  36. import builtins
  37. __all__ = [
  38. 'typename', 'is_tensor', 'is_storage',
  39. 'set_default_tensor_type', 'set_default_device', 'get_default_device',
  40. 'set_rng_state', 'get_rng_state', 'manual_seed', 'initial_seed', 'seed',
  41. 'save', 'load', 'set_printoptions', 'chunk', 'split', 'stack', 'matmul',
  42. 'no_grad', 'enable_grad', 'rand', 'randn', 'inference_mode',
  43. 'DoubleStorage', 'FloatStorage', 'LongStorage', 'IntStorage',
  44. 'ShortStorage', 'CharStorage', 'ByteStorage', 'BoolStorage',
  45. 'TypedStorage', 'UntypedStorage',
  46. 'DoubleTensor', 'FloatTensor', 'LongTensor', 'IntTensor',
  47. 'ShortTensor', 'CharTensor', 'ByteTensor', 'BoolTensor', 'Tensor',
  48. 'lobpcg', 'use_deterministic_algorithms',
  49. 'are_deterministic_algorithms_enabled',
  50. 'is_deterministic_algorithms_warn_only_enabled',
  51. 'set_deterministic_debug_mode', 'get_deterministic_debug_mode',
  52. 'set_float32_matmul_precision', 'get_float32_matmul_precision',
  53. 'set_warn_always', 'is_warn_always_enabled', 'SymInt', 'SymFloat',
  54. 'SymBool', 'sym_not', 'unravel_index',
  55. 'sym_int', 'sym_float', 'sym_max', 'sym_min', 'sym_ite', 'compile', 'vmap',
  56. 'export', 'autocast', 'cond', 'GradScaler',
  57. 'get_device_module',
  58. ]
  59. ################################################################################
  60. # Load the extension module
  61. ################################################################################
  62. if sys.platform == 'win32':
  63. import sysconfig
  64. pfiles_path = os.getenv('ProgramFiles', 'C:\\Program Files')
  65. py_dll_path = os.path.join(sys.exec_prefix, 'Library', 'bin')
  66. th_dll_path = os.path.join(os.path.dirname(__file__), 'lib')
  67. usebase_path = os.path.join(sysconfig.get_config_var("userbase"), 'Library', 'bin')
  68. # When users create a virtualenv that inherits the base environment,
  69. # we will need to add the corresponding library directory into
  70. # DLL search directories. Otherwise, it will rely on `PATH` which
  71. # is dependent on user settings.
  72. if sys.exec_prefix != sys.base_exec_prefix:
  73. base_py_dll_path = os.path.join(sys.base_exec_prefix, 'Library', 'bin')
  74. else:
  75. base_py_dll_path = ''
  76. dll_paths = list(filter(os.path.exists, [th_dll_path, py_dll_path, base_py_dll_path, usebase_path]))
  77. if all(not os.path.exists(os.path.join(p, 'nvToolsExt64_1.dll')) for p in dll_paths):
  78. nvtoolsext_dll_path = os.path.join(
  79. os.getenv('NVTOOLSEXT_PATH', os.path.join(pfiles_path, 'NVIDIA Corporation', 'NvToolsExt')), 'bin', 'x64')
  80. else:
  81. nvtoolsext_dll_path = ''
  82. from .version import cuda as cuda_version
  83. import glob
  84. if cuda_version and all(not glob.glob(os.path.join(p, 'cudart64*.dll')) for p in dll_paths):
  85. cuda_version_1 = cuda_version.replace('.', '_')
  86. cuda_path_var = 'CUDA_PATH_V' + cuda_version_1
  87. default_path = os.path.join(pfiles_path, 'NVIDIA GPU Computing Toolkit', 'CUDA', 'v' + cuda_version)
  88. cuda_path = os.path.join(os.getenv(cuda_path_var, default_path), 'bin')
  89. else:
  90. cuda_path = ''
  91. dll_paths.extend(filter(os.path.exists, [nvtoolsext_dll_path, cuda_path]))
  92. kernel32 = ctypes.WinDLL('kernel32.dll', use_last_error=True)
  93. with_load_library_flags = hasattr(kernel32, 'AddDllDirectory')
  94. prev_error_mode = kernel32.SetErrorMode(0x0001)
  95. kernel32.LoadLibraryW.restype = ctypes.c_void_p
  96. if with_load_library_flags:
  97. kernel32.LoadLibraryExW.restype = ctypes.c_void_p
  98. for dll_path in dll_paths:
  99. os.add_dll_directory(dll_path)
  100. try:
  101. ctypes.CDLL('vcruntime140.dll')
  102. ctypes.CDLL('msvcp140.dll')
  103. ctypes.CDLL('vcruntime140_1.dll')
  104. except OSError:
  105. print('''Microsoft Visual C++ Redistributable is not installed, this may lead to the DLL load failure.
  106. It can be downloaded at https://aka.ms/vs/16/release/vc_redist.x64.exe''')
  107. dlls = glob.glob(os.path.join(th_dll_path, '*.dll'))
  108. path_patched = False
  109. for dll in dlls:
  110. is_loaded = False
  111. if with_load_library_flags:
  112. res = kernel32.LoadLibraryExW(dll, None, 0x00001100)
  113. last_error = ctypes.get_last_error()
  114. if res is None and last_error != 126:
  115. err = ctypes.WinError(last_error)
  116. err.strerror += f' Error loading "{dll}" or one of its dependencies.'
  117. raise err
  118. elif res is not None:
  119. is_loaded = True
  120. if not is_loaded:
  121. if not path_patched:
  122. os.environ['PATH'] = ';'.join(dll_paths + [os.environ['PATH']])
  123. path_patched = True
  124. res = kernel32.LoadLibraryW(dll)
  125. if res is None:
  126. err = ctypes.WinError(ctypes.get_last_error())
  127. err.strerror += f' Error loading "{dll}" or one of its dependencies.'
  128. raise err
  129. kernel32.SetErrorMode(prev_error_mode)
  130. def _preload_cuda_deps(lib_folder, lib_name):
  131. """Preloads cuda deps if they could not be found otherwise."""
  132. # Should only be called on Linux if default path resolution have failed
  133. assert platform.system() == 'Linux', 'Should only be called on Linux'
  134. import glob
  135. lib_path = None
  136. for path in sys.path:
  137. nvidia_path = os.path.join(path, 'nvidia')
  138. if not os.path.exists(nvidia_path):
  139. continue
  140. candidate_lib_paths = glob.glob(os.path.join(nvidia_path, lib_folder, 'lib', lib_name))
  141. if candidate_lib_paths and not lib_path:
  142. lib_path = candidate_lib_paths[0]
  143. if lib_path:
  144. break
  145. if not lib_path:
  146. raise ValueError(f"{lib_name} not found in the system path {sys.path}")
  147. ctypes.CDLL(lib_path)
  148. # See Note [Global dependencies]
  149. def _load_global_deps() -> None:
  150. LIBTORCH_PKG_NAME = "libtorchsplit"
  151. def find_package_path(package_name):
  152. spec = importlib.util.find_spec(package_name)
  153. if spec:
  154. # The package might be a namespace package, so get_data may fail
  155. try:
  156. loader = spec.loader
  157. if loader is not None:
  158. file_path = loader.get_filename() # type: ignore[attr-defined]
  159. return os.path.dirname(file_path)
  160. except AttributeError:
  161. pass
  162. return None
  163. def load_shared_libraries(library_path):
  164. lib_dir = os.path.join(library_path, 'lib')
  165. if not os.path.exists(lib_dir):
  166. return
  167. # Determine the file extension based on the platform
  168. if platform.system() == 'Darwin':
  169. lib_ext = '.dylib'
  170. else:
  171. lib_ext = '.so'
  172. # Find all shared library files with the appropriate extension
  173. library_files = [f for f in os.listdir(lib_dir) if f.endswith(lib_ext)]
  174. if not library_files:
  175. return
  176. for lib_file in library_files:
  177. lib_path = os.path.join(lib_dir, lib_file)
  178. try:
  179. ctypes.CDLL(lib_path, mode=ctypes.RTLD_GLOBAL)
  180. except OSError as err:
  181. print(f"Failed to load {lib_path}: {err}")
  182. if _running_with_deploy() or platform.system() == 'Windows':
  183. return
  184. lib_name = 'libtorch_global_deps' + ('.dylib' if platform.system() == 'Darwin' else '.so')
  185. here = os.path.abspath(__file__)
  186. global_deps_lib_path = os.path.join(os.path.dirname(here), 'lib', lib_name)
  187. split_build_lib_name = LIBTORCH_PKG_NAME
  188. library_path = find_package_path(split_build_lib_name)
  189. if library_path:
  190. global_deps_lib_path = os.path.join(library_path, 'lib', lib_name)
  191. try:
  192. ctypes.CDLL(global_deps_lib_path, mode=ctypes.RTLD_GLOBAL)
  193. except OSError as err:
  194. # Can only happen for wheel with cuda libs as PYPI deps
  195. # As PyTorch is not purelib, but nvidia-*-cu12 is
  196. cuda_libs: Dict[str, str] = {
  197. 'cublas': 'libcublas.so.*[0-9]',
  198. 'cudnn': 'libcudnn.so.*[0-9]',
  199. 'cuda_nvrtc': 'libnvrtc.so.*[0-9]',
  200. 'cuda_runtime': 'libcudart.so.*[0-9]',
  201. 'cuda_cupti': 'libcupti.so.*[0-9]',
  202. 'cufft': 'libcufft.so.*[0-9]',
  203. 'curand': 'libcurand.so.*[0-9]',
  204. 'cusolver': 'libcusolver.so.*[0-9]',
  205. 'cusparse': 'libcusparse.so.*[0-9]',
  206. 'nccl': 'libnccl.so.*[0-9]',
  207. 'nvtx': 'libnvToolsExt.so.*[0-9]',
  208. }
  209. is_cuda_lib_err = [lib for lib in cuda_libs.values() if lib.split('.')[0] in err.args[0]]
  210. if not is_cuda_lib_err:
  211. raise err
  212. for lib_folder, lib_name in cuda_libs.items():
  213. _preload_cuda_deps(lib_folder, lib_name)
  214. ctypes.CDLL(global_deps_lib_path, mode=ctypes.RTLD_GLOBAL)
  215. if library_path:
  216. # loading libtorch_global_deps first due its special logic
  217. load_shared_libraries(library_path)
  218. if (USE_RTLD_GLOBAL_WITH_LIBTORCH or os.getenv('TORCH_USE_RTLD_GLOBAL')) and \
  219. (_running_with_deploy() or platform.system() != 'Windows'):
  220. # Do it the hard way. You might want to load libtorch with RTLD_GLOBAL in a
  221. # few circumstances:
  222. #
  223. # 1. You're in a build environment (e.g., fbcode) where
  224. # libtorch_global_deps is not available, but you still need
  225. # to get mkl to link in with RTLD_GLOBAL or it will just
  226. # not work.
  227. #
  228. # 2. You're trying to run PyTorch under UBSAN and you need
  229. # to ensure that only one copy of libtorch is loaded, so
  230. # vptr checks work properly
  231. #
  232. # If you're using this setting, you must verify that all the libraries
  233. # you load consistently use the same libstdc++, or you may have
  234. # mysterious segfaults.
  235. #
  236. old_flags = sys.getdlopenflags()
  237. sys.setdlopenflags(os.RTLD_GLOBAL | os.RTLD_LAZY)
  238. from torch._C import * # noqa: F403
  239. sys.setdlopenflags(old_flags)
  240. del old_flags
  241. else:
  242. # Easy way. You want this most of the time, because it will prevent
  243. # C++ symbols from libtorch clobbering C++ symbols from other
  244. # libraries, leading to mysterious segfaults.
  245. #
  246. # If building in an environment where libtorch_global_deps isn't available
  247. # like parts of fbsource, but where RTLD_GLOBAL causes segfaults, you will
  248. # want USE_RTLD_GLOBAL_WITH_LIBTORCH = False and USE_GLOBAL_DEPS = False
  249. #
  250. # See Note [Global dependencies]
  251. if USE_GLOBAL_DEPS:
  252. _load_global_deps()
  253. from torch._C import * # noqa: F403
  254. # Appease the type checker; ordinarily this binding is inserted by the
  255. # torch._C module initialization code in C
  256. if TYPE_CHECKING:
  257. from . import _C as _C # noqa: TCH004
  258. class SymInt:
  259. """
  260. Like an int (including magic methods), but redirects all operations on the
  261. wrapped node. This is used in particular to symbolically record operations
  262. in the symbolic shape workflow.
  263. """
  264. def __init__(self, node):
  265. # This field MUST be named node; C++ binding code assumes that this
  266. # class has a field named node that stores SymNode
  267. self.node = node
  268. def __bool__(self):
  269. return builtins.bool(self != 0)
  270. def __int__(self):
  271. return self.node.int_()
  272. def __index__(self):
  273. return self.node.int_()
  274. # Magic methods installed by torch.fx.experimental.sym_node
  275. def __round__(self, ndigits=None):
  276. return self
  277. def __truediv__(self, other):
  278. if isinstance(other, (builtins.float, SymFloat)):
  279. return sym_float(self).__float_truediv__(other)
  280. if not isinstance(other, (builtins.int, SymInt)):
  281. return NotImplemented
  282. return self.__int_truediv__(other)
  283. def __rtruediv__(self, other):
  284. if isinstance(other, (builtins.float, SymFloat)):
  285. return sym_float(self).__rfloat_truediv__(other)
  286. if not isinstance(other, (builtins.int, SymInt)):
  287. return NotImplemented
  288. return self.__rint_truediv__(other)
  289. def __floordiv__(self, other):
  290. if isinstance(other, (builtins.float, SymFloat)):
  291. return torch.sym_float(math.floor(sym_float(self) / other))
  292. if not isinstance(other, (builtins.int, SymInt)):
  293. return NotImplemented
  294. return self.__int_floordiv__(other)
  295. def __rfloordiv__(self, other):
  296. if isinstance(other, (builtins.float, SymFloat)):
  297. return torch.sym_float(math.floor(other / sym_float(self)))
  298. if not isinstance(other, (builtins.int, SymInt)):
  299. return NotImplemented
  300. return self.__rint_floordiv__(other)
  301. # nb: complex is impossible to handle correctly lol, with
  302. # negative base and integral float need to diverge semantics and
  303. # just always return complex. Neener neener pretend this problem
  304. # doesn't exist
  305. def __pow__(self, other):
  306. if isinstance(other, (builtins.float, SymFloat)):
  307. return sym_float(self).__pow__(other)
  308. if not isinstance(other, (builtins.int, SymInt)):
  309. return NotImplemented
  310. # Guards! This guard is necessary because we need to know it to
  311. # determine the output type of this operation
  312. if other >= 0:
  313. return self.__pow_by_natural__(other)
  314. else:
  315. # Mercifully, when the exponent is negative, Python just promotes
  316. # to doubles and does a float pow:
  317. #
  318. # if (Py_SIZE(b) < 0 && c == NULL) {
  319. # /* if exponent is negative and there's no modulus:
  320. # return a float. This works because we know
  321. # that this calls float_pow() which converts its
  322. # arguments to double. */
  323. # Py_DECREF(a);
  324. # Py_DECREF(b);
  325. # return PyFloat_Type.tp_as_number->nb_power(v, w, x);
  326. # }
  327. return sym_float(self).__pow__(sym_float(other))
  328. def __rpow__(self, other):
  329. if isinstance(other, (builtins.float, SymFloat)):
  330. return sym_float(self).__rpow__(other)
  331. if not isinstance(other, (builtins.int, SymInt)):
  332. return NotImplemented
  333. if self >= 0: # self is exponent
  334. return self.__rpow_by_natural__(other)
  335. else:
  336. return sym_float(self).__rpow__(sym_float(other))
  337. def __eq__(self, other: object) -> builtins.bool:
  338. raise AssertionError("type stub not overridden")
  339. def __lt__(self, other) -> builtins.bool:
  340. raise AssertionError("type stub not overridden")
  341. def __gt__(self, other) -> builtins.bool:
  342. raise AssertionError("type stub not overridden")
  343. def __le__(self, other) -> builtins.bool:
  344. raise AssertionError("type stub not overridden")
  345. def __ge__(self, other) -> builtins.bool:
  346. raise AssertionError("type stub not overridden")
  347. def __add__(self, other) -> "SymInt":
  348. raise AssertionError("type stub not overridden")
  349. def __mul__(self, other) -> "SymInt":
  350. raise AssertionError("type stub not overridden")
  351. def __pow_by_natural__(self, other) -> "SymInt":
  352. raise AssertionError("type stub not overridden")
  353. def __rpow_by_natural__(self, other) -> "SymInt":
  354. raise AssertionError("type stub not overridden")
  355. def __int_truediv__(self, other) -> "SymFloat":
  356. raise AssertionError("type stub not overridden")
  357. def __rint_truediv__(self, other) -> "SymFloat":
  358. raise AssertionError("type stub not overridden")
  359. def __int_floordiv__(self, other) -> "SymFloat":
  360. raise AssertionError("type stub not overridden")
  361. def __rint_floordiv__(self, other) -> "SymFloat":
  362. raise AssertionError("type stub not overridden")
  363. def __sym_max__(self, other):
  364. raise AssertionError("type stub not overridden")
  365. def __sym_min__(self, other):
  366. raise AssertionError("type stub not overridden")
  367. def __sym_float__(self):
  368. raise AssertionError("type stub not overridden")
  369. def __neg__(self):
  370. raise AssertionError("type stub not overridden")
  371. def __repr__(self):
  372. return str(self.node)
  373. def __hash__(self) -> builtins.int:
  374. if self.node.is_nested_int():
  375. return hash(self.node.nested_int())
  376. else:
  377. # We could support constant SymInts as well, but not doing it for now
  378. raise TypeError("unhashable type: non-nested SymInt")
  379. class SymFloat:
  380. """
  381. Like an float (including magic methods), but redirects all operations on the
  382. wrapped node. This is used in particular to symbolically record operations
  383. in the symbolic shape workflow.
  384. """
  385. def __init__(self, node):
  386. # This field MUST be named node; C++ binding code assumes that this
  387. # class has a field named node that stores SymNode
  388. self.node = node
  389. def __truediv__(self, other):
  390. if not isinstance(other, (builtins.int, builtins.float, SymInt, SymFloat)):
  391. return NotImplemented
  392. return self.__float_truediv__(sym_float(other))
  393. def __rtruediv__(self, other):
  394. if not isinstance(other, (builtins.int, builtins.float, SymInt, SymFloat)):
  395. return NotImplemented
  396. return self.__rfloat_truediv__(sym_float(other))
  397. def __floordiv__(self, other):
  398. if not isinstance(other, (builtins.int, builtins.float, SymInt, SymFloat)):
  399. return NotImplemented
  400. return torch.sym_float(math.floor(self / sym_float(other)))
  401. def __rfloordiv__(self, other):
  402. if not isinstance(other, (builtins.int, builtins.float, SymInt, SymFloat)):
  403. return NotImplemented
  404. return torch.sym_float(math.floor(sym_float(other) / self))
  405. def __bool__(self):
  406. return self.node.bool_()
  407. # Symbolic power does NOT work with negative base, this is to avoid
  408. # potential complex outputs
  409. def __pow__(self, other):
  410. if not isinstance(other, (builtins.int, builtins.float, SymInt, SymFloat)):
  411. return NotImplemented
  412. torch._check(self >= 0)
  413. return self.__float_pow__(other)
  414. def __rpow__(self, other):
  415. if not isinstance(other, (builtins.int, builtins.float, SymInt, SymFloat)):
  416. return NotImplemented
  417. torch._check(other >= 0)
  418. return self.__rfloat_pow__(other)
  419. # Magic methods installed by torch.fx.experimental.sym_node
  420. def __eq__(self, other: object) -> builtins.bool:
  421. raise AssertionError("type stub not overridden")
  422. def __lt__(self, other) -> builtins.bool:
  423. raise AssertionError("type stub not overridden")
  424. def __gt__(self, other) -> builtins.bool:
  425. raise AssertionError("type stub not overridden")
  426. def __le__(self, other) -> builtins.bool:
  427. raise AssertionError("type stub not overridden")
  428. def __ge__(self, other) -> builtins.bool:
  429. raise AssertionError("type stub not overridden")
  430. def __float_pow__(self, other) -> "SymFloat":
  431. raise AssertionError("type stub not overridden")
  432. def __rfloat_pow__(self, other) -> "SymFloat":
  433. raise AssertionError("type stub not overridden")
  434. def __float_truediv__(self, other) -> "SymFloat":
  435. raise AssertionError("type stub not overridden")
  436. def __rfloat_truediv__(self, other) -> "SymFloat":
  437. raise AssertionError("type stub not overridden")
  438. def __trunc__(self):
  439. raise AssertionError("type stub not overridden")
  440. def __sym_max__(self, other):
  441. raise AssertionError("type stub not overridden")
  442. def __sym_min__(self, other):
  443. raise AssertionError("type stub not overridden")
  444. def __sym_int__(self):
  445. raise AssertionError("type stub not overridden")
  446. def is_integer(self):
  447. """Return True if the float is an integer."""
  448. raise AssertionError("type stub not overridden")
  449. def __repr__(self):
  450. return self.node.str()
  451. class SymBool:
  452. """
  453. Like an bool (including magic methods), but redirects all operations on the
  454. wrapped node. This is used in particular to symbolically record operations
  455. in the symbolic shape workflow.
  456. Unlike regular bools, regular boolean operators will force extra guards instead
  457. of symbolically evaluate. Use the bitwise operators instead to handle this.
  458. """
  459. def __init__(self, node):
  460. # This field MUST be named node; C++ binding code assumes that this
  461. # class has a field named node that stores SymNode
  462. self.node = node
  463. def __bool__(self):
  464. return self.node.bool_()
  465. def __int__(self):
  466. return builtins.int(self.node.bool_())
  467. # Magic methods installed by torch.fx.experimental.sym_node
  468. def __and__(self, other) -> "SymBool":
  469. raise AssertionError("type stub not overridden")
  470. def __or__(self, other) -> "SymBool":
  471. raise AssertionError("type stub not overridden")
  472. # We very carefully define __sym_not__, and not a number of other
  473. # plausible alternatives:
  474. #
  475. # - We do not override __not__ because this is not a real magic
  476. # method; you cannot override the meaning of the not builtin in
  477. # Python. We use the name 'sym_not' to clarify that in user code you
  478. # cannot use the builtin not or operator.not_ or operator.__not__ and
  479. # hit this magic method; you must use our custom sym_not operator.
  480. #
  481. # - We do not override the __invert__ method because SymBool is
  482. # meant to be usable in situations where bool is expected. However,
  483. # bitwise negation ~a does the wrong thing with booleans (because
  484. # bool is a subclass of int, so ~1 = -2 which is not falseish.)
  485. # This would be a giant footgun, so we get around it by defining
  486. # our own operator. Note that bitwise and/or do the right thing,
  487. # so we reuse the conventional operators there for readability.
  488. #
  489. def __sym_not__(self) -> "SymBool":
  490. raise AssertionError("type stub not overridden")
  491. def __sym_ite__(self, then_val, else_val):
  492. raise AssertionError("type stub not overridden")
  493. def __eq__(self, other) -> builtins.bool:
  494. raise AssertionError("type stub not overridden")
  495. def __repr__(self):
  496. return str(self.node)
  497. def __hash__(self):
  498. if self.node.is_constant():
  499. return hash(self.node.bool_())
  500. else:
  501. raise TypeError("unhashable type: SymBool")
  502. def sym_not(a):
  503. r""" SymInt-aware utility for logical negation.
  504. Args:
  505. a (SymBool or bool): Object to negate
  506. """
  507. import sympy
  508. from .overrides import has_torch_function_unary, handle_torch_function
  509. if has_torch_function_unary(a):
  510. return handle_torch_function(sym_not, (a,), a)
  511. if hasattr(a, '__sym_not__'):
  512. return a.__sym_not__()
  513. if isinstance(a, sympy.Basic):
  514. return ~a # type: ignore[operator]
  515. return not a
  516. def sym_float(a):
  517. r""" SymInt-aware utility for float casting.
  518. Args:
  519. a (SymInt, SymFloat, or object): Object to cast
  520. """
  521. from .overrides import has_torch_function_unary, handle_torch_function
  522. if has_torch_function_unary(a):
  523. return handle_torch_function(sym_float, (a,), a)
  524. if isinstance(a, SymFloat):
  525. return a
  526. elif hasattr(a, '__sym_float__'):
  527. return a.__sym_float__()
  528. return py_float(a) # type: ignore[operator]
  529. def sym_int(a):
  530. r""" SymInt-aware utility for int casting.
  531. Args:
  532. a (SymInt, SymFloat, or object): Object to cast
  533. """
  534. from .overrides import has_torch_function_unary, handle_torch_function
  535. if has_torch_function_unary(a):
  536. return handle_torch_function(sym_int, (a,), a)
  537. if isinstance(a, SymInt):
  538. return a
  539. elif isinstance(a, SymFloat):
  540. return math.trunc(a)
  541. return py_int(a) # type: ignore[operator]
  542. def sym_max(a, b):
  543. """
  544. SymInt-aware utility for max which avoids branching on a < b.
  545. Unlike builtins.max(), this only works for int/float, and it always
  546. promotes to float if any argument is float (unlike builtins.max, which
  547. will faithfully preserve the type of the input argument).
  548. """
  549. from .overrides import has_torch_function, handle_torch_function
  550. if has_torch_function((a, b)):
  551. return handle_torch_function(sym_max, (a, b), a, b)
  552. if isinstance(a, (SymInt, SymFloat)):
  553. return a.__sym_max__(b)
  554. elif isinstance(b, (SymInt, SymFloat)):
  555. # Due to promotion semantics, this is operator is commutative:
  556. # max(1, 1.0) === max(1.0, 1) === 1.0
  557. return b.__sym_max__(a)
  558. # TODO: Probably can make bool work too, just lazy
  559. assert isinstance(a, (builtins.int, builtins.float)), type(a)
  560. assert isinstance(b, (builtins.int, builtins.float)), type(b)
  561. if isinstance(a, builtins.float) or isinstance(b, builtins.float):
  562. return builtins.float(builtins.max(a, b))
  563. else:
  564. return builtins.max(a, b)
  565. def sym_min(a, b):
  566. """ SymInt-aware utility for min()."""
  567. from .overrides import has_torch_function, handle_torch_function
  568. if has_torch_function((a, b)):
  569. return handle_torch_function(sym_min, (a, b), a, b)
  570. if isinstance(a, (SymInt, SymFloat)):
  571. return a.__sym_min__(b)
  572. elif isinstance(b, (SymInt, SymFloat)):
  573. return b.__sym_min__(a)
  574. assert isinstance(a, (builtins.int, builtins.float)), type(a)
  575. assert isinstance(b, (builtins.int, builtins.float)), type(b)
  576. if isinstance(a, builtins.float) or isinstance(b, builtins.float):
  577. return builtins.float(builtins.min(a, b))
  578. else:
  579. return builtins.min(a, b)
  580. # Drop in replacement for math.sqrt, math.sin, math.cos etc
  581. def _get_sym_math_fn(name):
  582. def fn(a):
  583. from .overrides import has_torch_function_unary, handle_torch_function
  584. if has_torch_function_unary(a):
  585. return handle_torch_function(fn, (a,), a)
  586. if hasattr(a, f"__sym_{name}__"):
  587. return getattr(a, f"__sym_{name}__")()
  588. return getattr(math, name)(a)
  589. return fn
  590. __fn, __name, __sym_name = None, '', ''
  591. for __name in ("sqrt", "cos", "cosh", "sin", "sinh", "tan", "tanh", "asin", "acos", "atan"):
  592. __sym_name = f"_sym_{__name}"
  593. __fn = _get_sym_math_fn(__name)
  594. __fn.__qualname__ = __fn.__name__ = __sym_name
  595. globals()[__sym_name] = __fn
  596. del __fn, __name, __sym_name, _get_sym_math_fn
  597. # Adding temporary shortcut
  598. sym_sqrt = globals()["_sym_sqrt"]
  599. __all__.append("sym_sqrt")
  600. def sym_ite(b, t, f):
  601. from .overrides import has_torch_function, handle_torch_function
  602. if has_torch_function((b, t, f)):
  603. return handle_torch_function(sym_ite, (b, t, f), b, t, f)
  604. assert isinstance(b, (SymBool, builtins.bool)) and type(t) == type(f)
  605. if isinstance(b, SymBool):
  606. return b.__sym_ite__(t, f)
  607. return t if b else f
  608. # Check to see if we can load C extensions, and if not provide some guidance
  609. # on what the problem might be.
  610. try:
  611. # _initExtension is chosen (arbitrarily) as a sentinel.
  612. from torch._C import _initExtension
  613. except ImportError:
  614. import torch._C as _C_for_compiled_check
  615. # The __file__ check only works for Python 3.7 and above.
  616. if _C_for_compiled_check.__file__ is None:
  617. raise ImportError(textwrap.dedent('''
  618. Failed to load PyTorch C extensions:
  619. It appears that PyTorch has loaded the `torch/_C` folder
  620. of the PyTorch repository rather than the C extensions which
  621. are expected in the `torch._C` namespace. This can occur when
  622. using the `install` workflow. e.g.
  623. $ python setup.py install && python -c "import torch"
  624. This error can generally be solved using the `develop` workflow
  625. $ python setup.py develop && python -c "import torch" # This should succeed
  626. or by running Python from a different directory.
  627. ''').strip()) from None
  628. raise # If __file__ is not None the cause is unknown, so just re-raise.
  629. __name, __obj = '', None
  630. for __name in dir(_C):
  631. if __name[0] != '_' and not __name.endswith('Base'):
  632. __all__.append(__name)
  633. __obj = getattr(_C, __name)
  634. if callable(__obj) or inspect.isclass(__obj):
  635. if __obj.__module__ != __name__:
  636. # TODO: fix their module from C++ side
  637. if __name not in ['DisableTorchFunctionSubclass', 'DisableTorchFunction', 'Generator']:
  638. __obj.__module__ = __name__
  639. elif __name == 'TensorBase':
  640. # issue 109438 / pr 109940. Prevent TensorBase from being copied into torch.
  641. delattr(sys.modules[__name__], __name)
  642. del __name, __obj
  643. if not TYPE_CHECKING:
  644. # issue 38137 and python issue 43367. Submodules of a C extension are
  645. # non-standard, and attributes of those submodules cannot be pickled since
  646. # pickle expect to be able to import them as "from _C.sub import attr"
  647. # which fails with "_C is not a package
  648. __name, __candidate = '', None
  649. for __name in dir(_C):
  650. __candidate = getattr(_C, __name)
  651. if type(__candidate) is type(_C):
  652. # submodule
  653. sys.modules.setdefault(f"{__name__}._C.{__name}", __candidate)
  654. del __name, __candidate
  655. ################################################################################
  656. # Define basic utilities
  657. ################################################################################
  658. def typename(o):
  659. if isinstance(o, torch.Tensor):
  660. return o.type()
  661. module = ''
  662. class_name = ''
  663. if hasattr(o, '__module__') and o.__module__ != 'builtins' \
  664. and o.__module__ != '__builtin__' and o.__module__ is not None:
  665. module = o.__module__ + '.'
  666. if hasattr(o, '__qualname__'):
  667. class_name = o.__qualname__
  668. elif hasattr(o, '__name__'):
  669. class_name = o.__name__
  670. else:
  671. class_name = o.__class__.__name__
  672. return module + class_name
  673. def is_tensor(obj):
  674. r"""Returns True if `obj` is a PyTorch tensor.
  675. Note that this function is simply doing ``isinstance(obj, Tensor)``.
  676. Using that ``isinstance`` check is better for typechecking with mypy,
  677. and more explicit - so it's recommended to use that instead of
  678. ``is_tensor``.
  679. Args:
  680. obj (Object): Object to test
  681. Example::
  682. >>> x = torch.tensor([1, 2, 3])
  683. >>> torch.is_tensor(x)
  684. True
  685. """
  686. return isinstance(obj, torch.Tensor)
  687. def is_storage(obj):
  688. r"""Returns True if `obj` is a PyTorch storage object.
  689. Args:
  690. obj (Object): Object to test
  691. """
  692. return type(obj) in _storage_classes
  693. _GLOBAL_DEVICE_CONTEXT = threading.local()
  694. def get_default_device() -> "torch.device":
  695. r"""Gets the default ``torch.Tensor`` to be allocated on ``device``"""
  696. global _GLOBAL_DEVICE_CONTEXT
  697. if hasattr(_GLOBAL_DEVICE_CONTEXT, "device_context"):
  698. device = _GLOBAL_DEVICE_CONTEXT.device_context.device
  699. if device.index is not None:
  700. return device
  701. else:
  702. # TODO: Call like get_device_index() method corresponding to
  703. # each device type
  704. return torch.tensor([]).device
  705. else:
  706. return torch.device("cpu")
  707. def set_default_device(device):
  708. """Sets the default ``torch.Tensor`` to be allocated on ``device``. This
  709. does not affect factory function calls which are called with an explicit
  710. ``device`` argument. Factory calls will be performed as if they
  711. were passed ``device`` as an argument.
  712. To only temporarily change the default device instead of setting it
  713. globally, use ``with torch.device(device):`` instead.
  714. The default device is initially ``cpu``. If you set the default tensor
  715. device to another device (e.g., ``cuda``) without a device index, tensors
  716. will be allocated on whatever the current device for the device type,
  717. even after :func:`torch.cuda.set_device` is called.
  718. .. warning::
  719. This function imposes a slight performance cost on every Python
  720. call to the torch API (not just factory functions). If this
  721. is causing problems for you, please comment on
  722. https://github.com/pytorch/pytorch/issues/92701
  723. .. note::
  724. This doesn't affect functions that create tensors that share the same memory as the input, like:
  725. :func:`torch.from_numpy` and :func:`torch.frombuffer`
  726. Args:
  727. device (device or string): the device to set as default
  728. Example::
  729. >>> # xdoctest: +SKIP("requires cuda, changes global state")
  730. >>> torch.get_default_device()
  731. device(type='cpu')
  732. >>> torch.set_default_device('cuda') # current device is 0
  733. >>> torch.get_default_device()
  734. device(type='cuda', index=0)
  735. >>> torch.set_default_device('cuda')
  736. >>> torch.cuda.set_device('cuda:1') # current device is 1
  737. >>> torch.get_default_device()
  738. device(type='cuda', index=1)
  739. >>> torch.set_default_device('cuda:1')
  740. >>> torch.get_default_device()
  741. device(type='cuda', index=1)
  742. """
  743. global _GLOBAL_DEVICE_CONTEXT
  744. if hasattr(_GLOBAL_DEVICE_CONTEXT, "device_context"):
  745. device_context = _GLOBAL_DEVICE_CONTEXT.device_context
  746. if device_context is not None:
  747. device_context.__exit__(None, None, None)
  748. if device is None:
  749. device_context = None
  750. else:
  751. from torch.utils._device import DeviceContext
  752. device_context = DeviceContext(device)
  753. device_context.__enter__()
  754. _GLOBAL_DEVICE_CONTEXT.device_context = device_context
  755. def set_default_tensor_type(t):
  756. r"""
  757. .. warning::
  758. This function is deprecated as of PyTorch 2.1, please use :func:`torch.set_default_dtype()` and
  759. :func:`torch.set_default_device()` as alternatives.
  760. Sets the default ``torch.Tensor`` type to floating point tensor type
  761. ``t``. This type will also be used as default floating point type for
  762. type inference in :func:`torch.tensor`.
  763. The default floating point tensor type is initially ``torch.FloatTensor``.
  764. Args:
  765. t (type or string): the floating point tensor type or its name
  766. Example::
  767. >>> # xdoctest: +SKIP("Other tests may have changed the default type. Can we reset it?")
  768. >>> torch.tensor([1.2, 3]).dtype # initial default for floating point is torch.float32
  769. torch.float32
  770. >>> torch.set_default_tensor_type(torch.DoubleTensor)
  771. >>> torch.tensor([1.2, 3]).dtype # a new floating point tensor
  772. torch.float64
  773. """
  774. if isinstance(t, str):
  775. t = _import_dotted_name(t)
  776. _C._set_default_tensor_type(t)
  777. def set_default_dtype(d):
  778. r"""
  779. Sets the default floating point dtype to :attr:`d`. Supports floating point dtype
  780. as inputs. Other dtypes will cause torch to raise an exception.
  781. When PyTorch is initialized its default floating point dtype is torch.float32,
  782. and the intent of set_default_dtype(torch.float64) is to facilitate NumPy-like
  783. type inference. The default floating point dtype is used to:
  784. 1. Implicitly determine the default complex dtype. When the default floating type is float16,
  785. the default complex dtype is complex32. For float32, the default complex dtype is complex64.
  786. For float64, it is complex128. For bfloat16, an exception will be raised because
  787. there is no corresponding complex type for bfloat16.
  788. 2. Infer the dtype for tensors constructed using Python floats or complex Python
  789. numbers. See examples below.
  790. 3. Determine the result of type promotion between bool and integer tensors and
  791. Python floats and complex Python numbers.
  792. Args:
  793. d (:class:`torch.dtype`): the floating point dtype to make the default.
  794. Either torch.float32 or torch.float64.
  795. Example:
  796. >>> # xdoctest: +SKIP("Other tests may have changed the default type. Can we reset it?")
  797. >>> # initial default for floating point is torch.float32
  798. >>> # Python floats are interpreted as float32
  799. >>> torch.tensor([1.2, 3]).dtype
  800. torch.float32
  801. >>> # initial default for floating point is torch.complex64
  802. >>> # Complex Python numbers are interpreted as complex64
  803. >>> torch.tensor([1.2, 3j]).dtype
  804. torch.complex64
  805. >>> torch.set_default_dtype(torch.float64)
  806. >>> # Python floats are now interpreted as float64
  807. >>> torch.tensor([1.2, 3]).dtype # a new floating point tensor
  808. torch.float64
  809. >>> # Complex Python numbers are now interpreted as complex128
  810. >>> torch.tensor([1.2, 3j]).dtype # a new complex tensor
  811. torch.complex128
  812. >>> torch.set_default_dtype(torch.float16)
  813. >>> # Python floats are now interpreted as float16
  814. >>> torch.tensor([1.2, 3]).dtype # a new floating point tensor
  815. torch.float16
  816. >>> # Complex Python numbers are now interpreted as complex128
  817. >>> torch.tensor([1.2, 3j]).dtype # a new complex tensor
  818. torch.complex32
  819. """
  820. _C._set_default_dtype(d)
  821. def use_deterministic_algorithms(mode: builtins.bool, *, warn_only: builtins.bool = False) -> None:
  822. r""" Sets whether PyTorch operations must use "deterministic"
  823. algorithms. That is, algorithms which, given the same input, and when
  824. run on the same software and hardware, always produce the same output.
  825. When enabled, operations will use deterministic algorithms when available,
  826. and if only nondeterministic algorithms are available they will throw a
  827. :class:`RuntimeError` when called.
  828. .. note:: This setting alone is not always enough to make an application
  829. reproducible. Refer to :ref:`reproducibility` for more information.
  830. .. note:: :func:`torch.set_deterministic_debug_mode` offers an alternative
  831. interface for this feature.
  832. The following normally-nondeterministic operations will act
  833. deterministically when ``mode=True``:
  834. * :class:`torch.nn.Conv1d` when called on CUDA tensor
  835. * :class:`torch.nn.Conv2d` when called on CUDA tensor
  836. * :class:`torch.nn.Conv3d` when called on CUDA tensor
  837. * :class:`torch.nn.ConvTranspose1d` when called on CUDA tensor
  838. * :class:`torch.nn.ConvTranspose2d` when called on CUDA tensor
  839. * :class:`torch.nn.ConvTranspose3d` when called on CUDA tensor
  840. * :class:`torch.nn.ReplicationPad2d` when attempting to differentiate a CUDA tensor
  841. * :func:`torch.bmm` when called on sparse-dense CUDA tensors
  842. * :func:`torch.Tensor.__getitem__` when attempting to differentiate a CPU tensor
  843. and the index is a list of tensors
  844. * :func:`torch.Tensor.index_put` with ``accumulate=False``
  845. * :func:`torch.Tensor.index_put` with ``accumulate=True`` when called on a CPU
  846. tensor
  847. * :func:`torch.Tensor.put_` with ``accumulate=True`` when called on a CPU
  848. tensor
  849. * :func:`torch.Tensor.scatter_add_` when called on a CUDA tensor
  850. * :func:`torch.gather` when called on a CUDA tensor that requires grad
  851. * :func:`torch.index_add` when called on CUDA tensor
  852. * :func:`torch.index_select` when attempting to differentiate a CUDA tensor
  853. * :func:`torch.repeat_interleave` when attempting to differentiate a CUDA tensor
  854. * :func:`torch.Tensor.index_copy` when called on a CPU or CUDA tensor
  855. * :func:`torch.Tensor.scatter` when `src` type is Tensor and called on CUDA tensor
  856. * :func:`torch.Tensor.scatter_reduce` when ``reduce='sum'`` or ``reduce='mean'`` and called on CUDA tensor
  857. The following normally-nondeterministic operations will throw a
  858. :class:`RuntimeError` when ``mode=True``:
  859. * :class:`torch.nn.AvgPool3d` when attempting to differentiate a CUDA tensor
  860. * :class:`torch.nn.AdaptiveAvgPool2d` when attempting to differentiate a CUDA tensor
  861. * :class:`torch.nn.AdaptiveAvgPool3d` when attempting to differentiate a CUDA tensor
  862. * :class:`torch.nn.MaxPool3d` when attempting to differentiate a CUDA tensor
  863. * :class:`torch.nn.AdaptiveMaxPool2d` when attempting to differentiate a CUDA tensor
  864. * :class:`torch.nn.FractionalMaxPool2d` when attempting to differentiate a CUDA tensor
  865. * :class:`torch.nn.FractionalMaxPool3d` when attempting to differentiate a CUDA tensor
  866. * :class:`torch.nn.MaxUnpool1d`
  867. * :class:`torch.nn.MaxUnpool2d`
  868. * :class:`torch.nn.MaxUnpool3d`
  869. * :func:`torch.nn.functional.interpolate` when attempting to differentiate a CUDA tensor
  870. and one of the following modes is used:
  871. - ``linear``
  872. - ``bilinear``
  873. - ``bicubic``
  874. - ``trilinear``
  875. * :class:`torch.nn.ReflectionPad1d` when attempting to differentiate a CUDA tensor
  876. * :class:`torch.nn.ReflectionPad2d` when attempting to differentiate a CUDA tensor
  877. * :class:`torch.nn.ReflectionPad3d` when attempting to differentiate a CUDA tensor
  878. * :class:`torch.nn.ReplicationPad1d` when attempting to differentiate a CUDA tensor
  879. * :class:`torch.nn.ReplicationPad3d` when attempting to differentiate a CUDA tensor
  880. * :class:`torch.nn.NLLLoss` when called on a CUDA tensor
  881. * :class:`torch.nn.CTCLoss` when attempting to differentiate a CUDA tensor
  882. * :class:`torch.nn.EmbeddingBag` when attempting to differentiate a CUDA tensor when
  883. ``mode='max'``
  884. * :func:`torch.Tensor.put_` when ``accumulate=False``
  885. * :func:`torch.Tensor.put_` when ``accumulate=True`` and called on a CUDA tensor
  886. * :func:`torch.histc` when called on a CUDA tensor
  887. * :func:`torch.bincount` when called on a CUDA tensor and ``weights``
  888. tensor is given
  889. * :func:`torch.kthvalue` with called on a CUDA tensor
  890. * :func:`torch.median` with indices output when called on a CUDA tensor
  891. * :func:`torch.nn.functional.grid_sample` when attempting to differentiate a CUDA tensor
  892. * :func:`torch.cumsum` when called on a CUDA tensor when dtype is floating point or complex
  893. * :func:`torch.Tensor.scatter_reduce` when ``reduce='prod'`` and called on CUDA tensor
  894. * :func:`torch.Tensor.resize_` when called with a quantized tensor
  895. In addition, several operations fill uninitialized memory when this setting
  896. is turned on and when
  897. :attr:`torch.utils.deterministic.fill_uninitialized_memory` is turned on.
  898. See the documentation for that attribute for more information.
  899. A handful of CUDA operations are nondeterministic if the CUDA version is
  900. 10.2 or greater, unless the environment variable ``CUBLAS_WORKSPACE_CONFIG=:4096:8``
  901. or ``CUBLAS_WORKSPACE_CONFIG=:16:8`` is set. See the CUDA documentation for more
  902. details: `<https://docs.nvidia.com/cuda/cublas/index.html#results-reproducibility>`_
  903. If one of these environment variable configurations is not set, a :class:`RuntimeError`
  904. will be raised from these operations when called with CUDA tensors:
  905. * :func:`torch.mm`
  906. * :func:`torch.mv`
  907. * :func:`torch.bmm`
  908. Note that deterministic operations tend to have worse performance than
  909. nondeterministic operations.
  910. .. note::
  911. This flag does not detect or prevent nondeterministic behavior caused
  912. by calling an inplace operation on a tensor with an internal memory
  913. overlap or by giving such a tensor as the :attr:`out` argument for an
  914. operation. In these cases, multiple writes of different data may target
  915. a single memory location, and the order of writes is not guaranteed.
  916. Args:
  917. mode (:class:`bool`): If True, makes potentially nondeterministic
  918. operations switch to a deterministic algorithm or throw a runtime
  919. error. If False, allows nondeterministic operations.
  920. Keyword args:
  921. warn_only (:class:`bool`, optional): If True, operations that do not
  922. have a deterministic implementation will throw a warning instead of
  923. an error. Default: ``False``
  924. Example::
  925. >>> # xdoctest: +SKIP
  926. >>> torch.use_deterministic_algorithms(True)
  927. # Forward mode nondeterministic error
  928. >>> torch.randn(10, device='cuda').kthvalue(1)
  929. ...
  930. RuntimeError: kthvalue CUDA does not have a deterministic implementation...
  931. # Backward mode nondeterministic error
  932. >>> torch.nn.AvgPool3d(1)(torch.randn(3, 4, 5, 6, requires_grad=True).cuda()).sum().backward()
  933. ...
  934. RuntimeError: avg_pool3d_backward_cuda does not have a deterministic implementation...
  935. """
  936. _C._set_deterministic_algorithms(mode, warn_only=warn_only)
  937. def are_deterministic_algorithms_enabled() -> builtins.bool:
  938. r"""Returns True if the global deterministic flag is turned on. Refer to
  939. :func:`torch.use_deterministic_algorithms` documentation for more details.
  940. """
  941. return _C._get_deterministic_algorithms()
  942. def is_deterministic_algorithms_warn_only_enabled() -> builtins.bool:
  943. r"""Returns True if the global deterministic flag is set to warn only.
  944. Refer to :func:`torch.use_deterministic_algorithms` documentation for more
  945. details.
  946. """
  947. return _C._get_deterministic_algorithms_warn_only()
  948. def set_deterministic_debug_mode(debug_mode: Union[builtins.int, str]) -> None:
  949. r"""Sets the debug mode for deterministic operations.
  950. .. note:: This is an alternative interface for
  951. :func:`torch.use_deterministic_algorithms`. Refer to that function's
  952. documentation for details about affected operations.
  953. Args:
  954. debug_mode(str or int): If "default" or 0, don't error or warn on
  955. nondeterministic operations. If "warn" or 1, warn on
  956. nondeterministic operations. If "error" or 2, error on
  957. nondeterministic operations.
  958. """
  959. # NOTE: builtins.int is used here because int in this scope resolves
  960. # to torch.int
  961. if not isinstance(debug_mode, (builtins.int, str)):
  962. raise TypeError(f'debug_mode must be str or int, but got {type(debug_mode)}')
  963. if isinstance(debug_mode, str):
  964. if debug_mode == 'default':
  965. debug_mode = 0
  966. elif debug_mode == 'warn':
  967. debug_mode = 1
  968. elif debug_mode == 'error':
  969. debug_mode = 2
  970. else:
  971. raise RuntimeError(
  972. 'invalid value of debug_mode, expected one of `default`, '
  973. f'`warn`, `error`, but got {debug_mode}')
  974. if debug_mode == 0:
  975. _C._set_deterministic_algorithms(False)
  976. elif debug_mode == 1:
  977. _C._set_deterministic_algorithms(True, warn_only=True)
  978. elif debug_mode == 2:
  979. _C._set_deterministic_algorithms(True)
  980. else:
  981. raise RuntimeError(
  982. 'invalid value of debug_mode, expected 0, 1, or 2, '
  983. f'but got {debug_mode}')
  984. def get_deterministic_debug_mode() -> builtins.int:
  985. r"""Returns the current value of the debug mode for deterministic
  986. operations. Refer to :func:`torch.set_deterministic_debug_mode`
  987. documentation for more details.
  988. """
  989. if _C._get_deterministic_algorithms():
  990. if _C._get_deterministic_algorithms_warn_only():
  991. return 1
  992. else:
  993. return 2
  994. else:
  995. return 0
  996. def get_float32_matmul_precision() -> builtins.str:
  997. r"""Returns the current value of float32 matrix multiplication precision. Refer to
  998. :func:`torch.set_float32_matmul_precision` documentation for more details.
  999. """
  1000. return _C._get_float32_matmul_precision()
  1001. def set_float32_matmul_precision(precision: str) -> None:
  1002. r"""Sets the internal precision of float32 matrix multiplications.
  1003. Running float32 matrix multiplications in lower precision may significantly increase
  1004. performance, and in some programs the loss of precision has a negligible impact.
  1005. Supports three settings:
  1006. * "highest", float32 matrix multiplications use the float32 datatype (24 mantissa
  1007. bits with 23 bits explicitly stored) for internal computations.
  1008. * "high", float32 matrix multiplications either use the TensorFloat32 datatype (10
  1009. mantissa bits explicitly stored) or treat each float32 number as the sum of two bfloat16 numbers
  1010. (approximately 16 mantissa bits with 14 bits explicitly stored), if the appropriate fast matrix multiplication
  1011. algorithms are available. Otherwise float32 matrix multiplications are computed
  1012. as if the precision is "highest". See below for more information on the bfloat16
  1013. approach.
  1014. * "medium", float32 matrix multiplications use the bfloat16 datatype (8 mantissa
  1015. bits with 7 bits explicitly stored) for internal computations, if a fast matrix multiplication algorithm
  1016. using that datatype internally is available. Otherwise float32
  1017. matrix multiplications are computed as if the precision is "high".
  1018. When using "high" precision, float32 multiplications may use a bfloat16-based algorithm
  1019. that is more complicated than simply truncating to some smaller number mantissa bits
  1020. (e.g. 10 for TensorFloat32, 7 for bfloat16 explicitly stored). Refer to [Henry2019]_ for a complete
  1021. description of this algorithm. To briefly explain here, the first step is to realize
  1022. that we can perfectly encode a single float32 number as the sum of three bfloat16
  1023. numbers (because float32 has 23 mantissa bits while bfloat16 has 7 explicitly stored, and both have the
  1024. same number of exponent bits). This means that the product of two float32 numbers can
  1025. be exactly given by the sum of nine products of bfloat16 numbers. We can then trade
  1026. accuracy for speed by dropping some of these products. The "high" precision algorithm
  1027. specifically keeps only the three most significant products, which conveniently excludes
  1028. all of the products involving the last 8 mantissa bits of either input. This means that
  1029. we can represent our inputs as the sum of two bfloat16 numbers rather than three.
  1030. Because bfloat16 fused-multiply-add (FMA) instructions are typically >10x faster than
  1031. float32 ones, it's faster to do three multiplications and 2 additions with bfloat16
  1032. precision than it is to do a single multiplication with float32 precision.
  1033. .. [Henry2019] http://arxiv.org/abs/1904.06376
  1034. .. note::
  1035. This does not change the output dtype of float32 matrix multiplications,
  1036. it controls how the internal computation of the matrix multiplication is performed.
  1037. .. note::
  1038. This does not change the precision of convolution operations. Other flags,
  1039. like `torch.backends.cudnn.allow_tf32`, may control the precision of convolution
  1040. operations.
  1041. .. note::
  1042. This flag currently only affects one native device type: CUDA.
  1043. If "high" or "medium" are set then the TensorFloat32 datatype will be used
  1044. when computing float32 matrix multiplications, equivalent to setting
  1045. `torch.backends.cuda.matmul.allow_tf32 = True`. When "highest" (the default)
  1046. is set then the float32 datatype is used for internal computations, equivalent
  1047. to setting `torch.backends.cuda.matmul.allow_tf32 = False`.
  1048. Args:
  1049. precision(str): can be set to "highest" (default), "high", or "medium" (see above).
  1050. """
  1051. _C._set_float32_matmul_precision(precision)
  1052. def set_warn_always(b: builtins.bool) -> None:
  1053. r"""When this flag is False (default) then some PyTorch warnings may only
  1054. appear once per process. This helps avoid excessive warning information.
  1055. Setting it to True causes these warnings to always appear, which may be
  1056. helpful when debugging.
  1057. Args:
  1058. b (:class:`bool`): If True, force warnings to always be emitted
  1059. If False, set to the default behaviour
  1060. """
  1061. _C._set_warnAlways(b)
  1062. def is_warn_always_enabled() -> builtins.bool:
  1063. r"""Returns True if the global warn_always flag is turned on. Refer to
  1064. :func:`torch.set_warn_always` documentation for more details.
  1065. """
  1066. return _C._get_warnAlways()
  1067. ################################################################################
  1068. # Define error checking functions
  1069. ################################################################################
  1070. # These error checking functions must be kept consistent with their C++
  1071. # equivalents. Their C++ equivalents are mentioned where applicable.
  1072. def _check_with(error_type, cond: Union[builtins.bool, SymBool], message: Callable[[], str]): # noqa: F811
  1073. if not isinstance(cond, (builtins.bool, torch.SymBool)):
  1074. raise TypeError(f'cond must be a bool, but got {type(cond)}')
  1075. from torch.fx.experimental.symbolic_shapes import expect_true
  1076. if expect_true(cond):
  1077. return
  1078. # error_type must be a subclass of Exception and not subclass of Warning
  1079. assert issubclass(error_type, Exception) and not issubclass(error_type, Warning)
  1080. if message is None:
  1081. message_evaluated = (
  1082. 'Expected cond to be True, but got False. (Could this error '
  1083. 'message be improved? If so, please report an enhancement request '
  1084. 'to PyTorch.)')
  1085. else:
  1086. if not callable(message):
  1087. raise TypeError('message must be a callable')
  1088. message_evaluated = str(message())
  1089. raise error_type(message_evaluated)
  1090. def _check(cond, message=None): # noqa: F811
  1091. r"""Throws error containing an optional message if the specified condition
  1092. is False.
  1093. Error type: ``RuntimeError``
  1094. C++ equivalent: ``TORCH_CHECK``
  1095. Args:
  1096. cond (:class:`bool`): If False, throw error
  1097. message (Callable, optional): Callable that returns either a string or
  1098. an object that has a ``__str__()`` method to be used as the error
  1099. message. Default: ``None``
  1100. """
  1101. _check_with(RuntimeError, cond, message)
  1102. def _check_is_size(i, message=None):
  1103. """Checks that a given integer is a valid size (i.e., is non-negative).
  1104. You should use this over _check(i >= 0) because we can use the semantic
  1105. information (that i is a size) to make some further inferences in case
  1106. i is an unbacked SymInt.
  1107. NB: Do NOT use this in contexts where a -1 size would be valid (indicating
  1108. to infer the size from context, or if you should wrap-around or truncate).
  1109. Only use this if the only valid value is an honest to goodness size.
  1110. """
  1111. # This is responsible for the expect_true
  1112. _check(i >= 0, message)
  1113. from torch.fx.experimental.symbolic_shapes import _advise_is_size
  1114. _advise_is_size(i)
  1115. def _check_index(cond, message=None): # noqa: F811
  1116. r"""Throws error containing an optional message if the specified condition
  1117. is False.
  1118. Error type: ``IndexError``
  1119. C++ equivalent: ``TORCH_CHECK_INDEX``
  1120. Args:
  1121. cond (:class:`bool`): If False, throw error
  1122. message (Callable, optional): Callable that returns either a string or
  1123. an object that has a ``__str__()`` method to be used as the error
  1124. message. Default: ``None``
  1125. """
  1126. _check_with(IndexError, cond, message)
  1127. def _check_value(cond, message=None): # noqa: F811
  1128. r"""Throws error containing an optional message if the specified condition
  1129. is False.
  1130. Error type: ``ValueError``
  1131. C++ equivalent: ``TORCH_CHECK_VALUE``
  1132. Args:
  1133. cond (:class:`bool`): If False, throw error
  1134. message (Callable, optional): Callable that returns either a string or
  1135. an object that has a ``__str__()`` method to be used as the error
  1136. message. Default: ``None``
  1137. """
  1138. _check_with(ValueError, cond, message)
  1139. def _check_type(cond, message=None): # noqa: F811
  1140. r"""Throws error containing an optional message if the specified condition
  1141. is False.
  1142. Error type: ``TypeError``
  1143. C++ equivalent: ``TORCH_CHECK_TYPE``
  1144. Args:
  1145. cond (:class:`bool`): If False, throw error
  1146. message (Callable, optional): Callable that returns either a string or
  1147. an object that has a ``__str__()`` method to be used as the error
  1148. message. Default: ``None``
  1149. """
  1150. _check_with(TypeError, cond, message)
  1151. def _check_not_implemented(cond, message=None): # noqa: F811
  1152. r"""Throws error containing an optional message if the specified condition
  1153. is False.
  1154. Error type: ``NotImplementedError``
  1155. C++ equivalent: ``TORCH_CHECK_NOT_IMPLEMENTED``
  1156. Args:
  1157. cond (:class:`bool`): If False, throw error
  1158. message (Callable, optional): Callable that returns either a string or
  1159. an object that has a ``__str__()`` method to be used as the error
  1160. message. Default: ``None``
  1161. """
  1162. _check_with(NotImplementedError, cond, message)
  1163. def _check_tensor_all_with(error_type, cond, message=None): # noqa: F811
  1164. if not torch.is_tensor(cond):
  1165. raise TypeError(f'cond must be a tensor, but got {type(cond)}')
  1166. if not cond.dtype == torch.bool:
  1167. raise TypeError(
  1168. f'cond tensor must have dtype torch.bool, but got {cond.dtype}')
  1169. _check_with(error_type, cond._is_all_true().item(), message)
  1170. # C++ equivalent: `TORCH_CHECK_TENSOR_ALL`
  1171. def _check_tensor_all(cond, message=None): # noqa: F811
  1172. r"""Throws error containing an optional message if the specified condition
  1173. is False.
  1174. Error type: ``RuntimeError``
  1175. C++ equivalent: ``TORCH_CHECK_TENSOR_ALL``
  1176. Args:
  1177. cond (:class:`torch.Tensor`): Tensor of dtype ``torch.bool``. If any
  1178. element is ``False``, throw error
  1179. message (Callable, optional): Callable that returns either a string or
  1180. an object that has a ``__str__()`` method to be used as the error
  1181. message. Default: ``None``
  1182. """
  1183. _check_tensor_all_with(RuntimeError, cond, message)
  1184. ################################################################################
  1185. # Define numeric constants
  1186. ################################################################################
  1187. # For Python Array API (https://data-apis.org/array-api/latest/API_specification/constants.html) and
  1188. # NumPy consistency (https://numpy.org/devdocs/reference/constants.html)
  1189. from math import e, nan , inf , pi
  1190. newaxis: None = None
  1191. __all__.extend(['e', 'pi', 'nan', 'inf', 'newaxis'])
  1192. ################################################################################
  1193. # Define Storage and Tensor classes
  1194. ################################################################################
  1195. from ._tensor import Tensor
  1196. from .storage import _StorageBase, TypedStorage, _LegacyStorage, UntypedStorage, _warn_typed_storage_removal
  1197. # NOTE: New <type>Storage classes should never be added. When adding a new
  1198. # dtype, use torch.storage.TypedStorage directly.
  1199. class ByteStorage(_LegacyStorage):
  1200. @classproperty
  1201. def dtype(self):
  1202. _warn_typed_storage_removal(stacklevel=3)
  1203. return self._dtype
  1204. @classproperty
  1205. def _dtype(self):
  1206. return torch.uint8
  1207. class DoubleStorage(_LegacyStorage):
  1208. @classproperty
  1209. def dtype(self):
  1210. _warn_typed_storage_removal(stacklevel=3)
  1211. return self._dtype
  1212. @classproperty
  1213. def _dtype(self):
  1214. return torch.double
  1215. class FloatStorage(_LegacyStorage):
  1216. @classproperty
  1217. def dtype(self):
  1218. _warn_typed_storage_removal(stacklevel=3)
  1219. return self._dtype
  1220. @classproperty
  1221. def _dtype(self):
  1222. return torch.float
  1223. class HalfStorage(_LegacyStorage):
  1224. @classproperty
  1225. def dtype(self):
  1226. _warn_typed_storage_removal(stacklevel=3)
  1227. return self._dtype
  1228. @classproperty
  1229. def _dtype(self):
  1230. return torch.half
  1231. class LongStorage(_LegacyStorage):
  1232. @classproperty
  1233. def dtype(self):
  1234. _warn_typed_storage_removal(stacklevel=3)
  1235. return self._dtype
  1236. @classproperty
  1237. def _dtype(self):
  1238. return torch.long
  1239. class IntStorage(_LegacyStorage):
  1240. @classproperty
  1241. def dtype(self):
  1242. _warn_typed_storage_removal(stacklevel=3)
  1243. return self._dtype
  1244. @classproperty
  1245. def _dtype(self):
  1246. return torch.int
  1247. class ShortStorage(_LegacyStorage):
  1248. @classproperty
  1249. def dtype(self):
  1250. _warn_typed_storage_removal(stacklevel=3)
  1251. return self._dtype
  1252. @classproperty
  1253. def _dtype(self):
  1254. return torch.short
  1255. class CharStorage(_LegacyStorage):
  1256. @classproperty
  1257. def dtype(self):
  1258. _warn_typed_storage_removal(stacklevel=3)
  1259. return self._dtype
  1260. @classproperty
  1261. def _dtype(self):
  1262. return torch.int8
  1263. class BoolStorage(_LegacyStorage):
  1264. @classproperty
  1265. def dtype(self):
  1266. _warn_typed_storage_removal(stacklevel=3)
  1267. return self._dtype
  1268. @classproperty
  1269. def _dtype(self):
  1270. return torch.bool
  1271. class BFloat16Storage(_LegacyStorage):
  1272. @classproperty
  1273. def dtype(self):
  1274. _warn_typed_storage_removal(stacklevel=3)
  1275. return self._dtype
  1276. @classproperty
  1277. def _dtype(self):
  1278. return torch.bfloat16
  1279. class ComplexDoubleStorage(_LegacyStorage):
  1280. @classproperty
  1281. def dtype(self):
  1282. _warn_typed_storage_removal(stacklevel=3)
  1283. return self._dtype
  1284. @classproperty
  1285. def _dtype(self):
  1286. return torch.cdouble
  1287. class ComplexFloatStorage(_LegacyStorage):
  1288. @classproperty
  1289. def dtype(self):
  1290. _warn_typed_storage_removal(stacklevel=3)
  1291. return self._dtype
  1292. @classproperty
  1293. def _dtype(self):
  1294. return torch.cfloat
  1295. class QUInt8Storage(_LegacyStorage):
  1296. @classproperty
  1297. def dtype(self):
  1298. _warn_typed_storage_removal(stacklevel=3)
  1299. return self._dtype
  1300. @classproperty
  1301. def _dtype(self):
  1302. return torch.quint8
  1303. class QInt8Storage(_LegacyStorage):
  1304. @classproperty
  1305. def dtype(self):
  1306. _warn_typed_storage_removal(stacklevel=3)
  1307. return self._dtype
  1308. @classproperty
  1309. def _dtype(self):
  1310. return torch.qint8
  1311. class QInt32Storage(_LegacyStorage):
  1312. @classproperty
  1313. def dtype(self):
  1314. _warn_typed_storage_removal(stacklevel=3)
  1315. return self._dtype
  1316. @classproperty
  1317. def _dtype(self):
  1318. return torch.qint32
  1319. class QUInt4x2Storage(_LegacyStorage):
  1320. @classproperty
  1321. def dtype(self):
  1322. _warn_typed_storage_removal(stacklevel=3)
  1323. return self._dtype
  1324. @classproperty
  1325. def _dtype(self):
  1326. return torch.quint4x2
  1327. class QUInt2x4Storage(_LegacyStorage):
  1328. @classproperty
  1329. def dtype(self):
  1330. _warn_typed_storage_removal(stacklevel=3)
  1331. return self._dtype
  1332. @classproperty
  1333. def _dtype(self):
  1334. return torch.quint2x4
  1335. _storage_classes = {
  1336. UntypedStorage, DoubleStorage, FloatStorage, LongStorage, IntStorage,
  1337. ShortStorage, CharStorage, ByteStorage, HalfStorage, BoolStorage,
  1338. QUInt8Storage, QInt8Storage, QInt32Storage, BFloat16Storage,
  1339. ComplexFloatStorage, ComplexDoubleStorage, QUInt4x2Storage, QUInt2x4Storage,
  1340. TypedStorage
  1341. }
  1342. # The _tensor_classes set is initialized by the call to initialize_python_bindings.
  1343. _tensor_classes: Set[Type] = set()
  1344. # If you edit these imports, please update torch/__init__.py.in as well
  1345. from .random import set_rng_state, get_rng_state, manual_seed, initial_seed, seed
  1346. from .serialization import save, load
  1347. from ._tensor_str import set_printoptions
  1348. ################################################################################
  1349. # Initialize extension
  1350. ################################################################################
  1351. def _manager_path():
  1352. if _running_with_deploy() or platform.system() == 'Windows':
  1353. return b""
  1354. path = get_file_path('torch', 'bin', 'torch_shm_manager')
  1355. prepare_multiprocessing_environment(get_file_path('torch'))
  1356. if not os.path.exists(path):
  1357. raise RuntimeError("Unable to find torch_shm_manager at " + path)
  1358. return path.encode('utf-8')
  1359. from torch.amp import autocast, GradScaler
  1360. # Initializing the extension shadows the built-in python float / int classes;
  1361. # store them for later use by SymInt / SymFloat.
  1362. py_float = float
  1363. py_int = int
  1364. # Shared memory manager needs to know the exact location of manager executable
  1365. _C._initExtension(_manager_path())
  1366. del _manager_path
  1367. # Appease the type checker: it can't deal with direct setting of globals().
  1368. # Note that we will see "too many" functions when reexporting this way; there
  1369. # is not a good way to fix this problem. Perhaps, try to redesign VariableFunctions
  1370. # so that this import is good enough
  1371. if TYPE_CHECKING:
  1372. # Some type signatures pulled in from _VariableFunctions here clash with
  1373. # signatures already imported. For now these clashes are ignored; see
  1374. # PR #43339 for details.
  1375. from torch._C._VariableFunctions import * # type: ignore[assignment, misc] # noqa: F403
  1376. # Fixup segment_reduce visibility
  1377. _segment_reduce = segment_reduce
  1378. del segment_reduce # noqa: F821
  1379. # Ops not to be exposed in `torch` namespace,
  1380. # mostly helper ops.
  1381. PRIVATE_OPS = (
  1382. 'unique_dim',
  1383. )
  1384. __name, __obj = '', None
  1385. for __name in dir(_C._VariableFunctions):
  1386. if __name.startswith('__') or __name in PRIVATE_OPS:
  1387. continue
  1388. __obj = getattr(_C._VariableFunctions, __name)
  1389. __obj.__module__ = __name__
  1390. # Hide some APIs that should not be public
  1391. if __name == "segment_reduce":
  1392. # TODO: Once the undocumented FC window is passed, remove the line bellow
  1393. globals()[__name] = __obj
  1394. __name = "_" + __name
  1395. globals()[__name] = __obj
  1396. if not __name.startswith("_"):
  1397. __all__.append(__name)
  1398. del __name, __obj
  1399. ################################################################################
  1400. # Add torch.dtype instances to the public API
  1401. ################################################################################
  1402. import torch
  1403. __all__.extend(
  1404. name for name in dir(torch) if isinstance(getattr(torch, name), torch.dtype)
  1405. )
  1406. ################################################################################
  1407. # Import TorchDynamo's lazy APIs to avoid circular dependenices
  1408. ################################################################################
  1409. # needs to be before from .functional import * to avoid circular dependencies
  1410. from ._compile import _disable_dynamo
  1411. ################################################################################
  1412. # Import interface functions defined in Python
  1413. ################################################################################
  1414. # needs to be after the above ATen bindings so we can overwrite from Python side
  1415. from .functional import * # noqa: F403
  1416. ################################################################################
  1417. # Remove unnecessary members
  1418. ################################################################################
  1419. del _StorageBase
  1420. del _LegacyStorage
  1421. ################################################################################
  1422. # Define _assert
  1423. ################################################################################
  1424. # needs to be before the submodule imports to avoid circular dependencies
  1425. def _assert(condition, message):
  1426. r"""A wrapper around Python's assert which is symbolically traceable.
  1427. """
  1428. from .overrides import has_torch_function, handle_torch_function
  1429. if type(condition) is not torch.Tensor and has_torch_function((condition,)):
  1430. return handle_torch_function(_assert, (condition,), condition, message)
  1431. assert condition, message
  1432. ################################################################################
  1433. # Import most common subpackages
  1434. ################################################################################
  1435. # Use the redundant form so that type checkers know that these are a part of
  1436. # the public API. The "regular" import lines are there solely for the runtime
  1437. # side effect of adding to the imported module's members for other users.
  1438. from torch import cuda as cuda
  1439. from torch import cpu as cpu
  1440. from torch import mps as mps
  1441. from torch import xpu as xpu
  1442. from torch import mtia as mtia
  1443. from torch import autograd as autograd
  1444. from torch.autograd import (
  1445. no_grad as no_grad,
  1446. enable_grad as enable_grad,
  1447. set_grad_enabled as set_grad_enabled,
  1448. inference_mode as inference_mode,
  1449. )
  1450. from torch import fft as fft
  1451. from torch import futures as futures
  1452. from torch import _awaits as _awaits
  1453. from torch import nested as nested
  1454. from torch import nn as nn
  1455. from torch.signal import windows as windows
  1456. from torch import optim as optim
  1457. import torch.optim._multi_tensor
  1458. from torch import multiprocessing as multiprocessing
  1459. from torch import sparse as sparse
  1460. from torch import special as special
  1461. import torch.utils.backcompat
  1462. from torch import jit as jit
  1463. from torch import linalg as linalg
  1464. from torch import hub as hub
  1465. from torch import random as random
  1466. from torch import distributions as distributions
  1467. from torch import testing as testing
  1468. from torch import backends as backends
  1469. import torch.utils.data
  1470. from torch import __config__ as __config__
  1471. from torch import __future__ as __future__
  1472. from torch import profiler as profiler
  1473. # Quantized, sparse, AO, etc. should be last to get imported, as nothing
  1474. # is expected to depend on them.
  1475. from torch import ao as ao
  1476. # nn.quant* depends on ao -- so should be after those.
  1477. import torch.nn.quantizable
  1478. import torch.nn.quantized
  1479. import torch.nn.qat
  1480. import torch.nn.intrinsic
  1481. _C._init_names(list(torch._storage_classes))
  1482. # attach docstrings to torch and tensor functions
  1483. from . import _torch_docs, _tensor_docs, _storage_docs, _size_docs
  1484. del _torch_docs, _tensor_docs, _storage_docs, _size_docs
  1485. def compiled_with_cxx11_abi() -> builtins.bool:
  1486. r"""Returns whether PyTorch was built with _GLIBCXX_USE_CXX11_ABI=1"""
  1487. return _C._GLIBCXX_USE_CXX11_ABI
  1488. # Import the ops "namespace"
  1489. from torch._ops import ops
  1490. from torch._classes import classes
  1491. import torch._library
  1492. # quantization depends on torch.fx
  1493. # Import quantization
  1494. from torch import quantization as quantization
  1495. # Import the quasi random sampler
  1496. from torch import quasirandom as quasirandom
  1497. # If you are seeing this, it means that this call site was not checked if
  1498. # the memory format could be preserved, and it was switched to old default
  1499. # behaviour of contiguous
  1500. legacy_contiguous_format = contiguous_format
  1501. # Register fork handler to initialize OpenMP in child processes (see gh-28389)
  1502. from torch.multiprocessing._atfork import register_after_fork
  1503. register_after_fork(torch.get_num_threads)
  1504. del register_after_fork
  1505. # Import tools that require fully imported torch (for applying
  1506. # torch.jit.script as a decorator, for instance):
  1507. from ._lobpcg import lobpcg as lobpcg
  1508. # These were previously defined in native_functions.yaml and appeared on the
  1509. # `torch` namespace, but we moved them to c10 dispatch to facilitate custom
  1510. # class usage. We add these lines here to preserve backward compatibility.
  1511. quantized_lstm = torch.ops.aten.quantized_lstm
  1512. quantized_gru = torch.ops.aten.quantized_gru
  1513. from torch.utils.dlpack import from_dlpack, to_dlpack
  1514. # Import experimental masked operations support. See
  1515. # [RFC-0016](https://github.com/pytorch/rfcs/pull/27) for more
  1516. # information.
  1517. from . import masked
  1518. # Import removed ops with error message about removal
  1519. from ._linalg_utils import ( # type: ignore[misc]
  1520. matrix_rank,
  1521. eig,
  1522. solve,
  1523. lstsq,
  1524. )
  1525. from ._linalg_utils import _symeig as symeig # type: ignore[misc]
  1526. class _TorchCompileInductorWrapper:
  1527. compiler_name = "inductor"
  1528. def __init__(self, mode, options, dynamic):
  1529. self.config: Dict[str, Any] = dict()
  1530. self.dynamic = dynamic
  1531. self.apply_mode(mode)
  1532. self.apply_options(options)
  1533. # Stash the compiler_fn to be used for backend match guard.
  1534. from torch._inductor.compile_fx import compile_fx
  1535. self.compiler_fn = compile_fx
  1536. if self.config.get("triton.cudagraphs", False):
  1537. os.environ["DISABLE_CUPTI_LAZY_REINIT"] = "1"
  1538. # FIXME: CUDA Graph does not work well with CUPTI teardown.
  1539. # 1) crashes on 1st lazy CUPTI re-init after teardown (CUDA 11)
  1540. # 2) crashes on 2nd non-lazy CUPTI re-init after teardown (CUDA 12)
  1541. # Workaround: turn off CUPTI teardown when using CUDA Graphs.
  1542. os.environ["TEARDOWN_CUPTI"] = "0"
  1543. def __eq__(self, other):
  1544. return (isinstance(other, _TorchCompileInductorWrapper) and
  1545. self.config == other.config and
  1546. self.dynamic == other.dynamic)
  1547. def apply_mode(self, mode: Optional[str]):
  1548. if mode is None or mode == "default":
  1549. pass
  1550. elif mode in ("reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"):
  1551. from torch._inductor import list_mode_options
  1552. self.apply_options(list_mode_options(mode, self.dynamic))
  1553. else:
  1554. raise RuntimeError(
  1555. f"Unrecognized mode={mode}, should be one of: default, reduce-overhead, max-autotune, max-autotune-no-cudagraphs"
  1556. )
  1557. def apply_options(self, options: Optional[Dict[str, Any]]):
  1558. if not options:
  1559. return
  1560. from torch._inductor import config
  1561. current_config: Dict[str, Any] = config.shallow_copy_dict()
  1562. for key, val in options.items():
  1563. attr_name = key.replace("-", "_")
  1564. if attr_name not in current_config:
  1565. raise RuntimeError(
  1566. f"Unexpected optimization option {key}, known options are {list(current_config.keys())}"
  1567. )
  1568. if type(val) is not type(current_config[attr_name]):
  1569. val_type_str = type(val).__name__
  1570. expected_type_str = type(current_config[attr_name]).__name__
  1571. raise RuntimeError(
  1572. f"Unexpected type of attr {key}, got {val_type_str} should be {expected_type_str}"
  1573. )
  1574. self.config[attr_name] = val
  1575. def __call__(self, model_, inputs_):
  1576. from torch._inductor.compile_fx import compile_fx
  1577. return compile_fx(model_, inputs_, config_patches=self.config)
  1578. def get_compiler_config(self):
  1579. from torch._inductor.compile_fx import get_patched_config_dict
  1580. return get_patched_config_dict(config_patches=self.config)
  1581. def reset(self):
  1582. from torch._inductor import config
  1583. if "triton.cudagraphs" in self.config or config.triton.cudagraphs:
  1584. if self.config.get("triton.cudagraphs", True):
  1585. from torch._inductor.cudagraph_trees import reset_cudagraph_trees
  1586. reset_cudagraph_trees()
  1587. class _TorchCompileWrapper:
  1588. def __init__(self, backend, mode, options, dynamic):
  1589. from torch._dynamo.backends.registry import lookup_backend
  1590. if isinstance(backend, str):
  1591. self.compiler_name = backend
  1592. elif hasattr(backend, "__name__"):
  1593. self.compiler_name = backend.__name__
  1594. else:
  1595. self.compiler_name = str(backend)
  1596. self.dynamic = dynamic
  1597. self.compiler_fn = lookup_backend(backend)
  1598. self.kwargs = {}
  1599. # only pass the args if they non-empty
  1600. if mode and mode != "default":
  1601. self.kwargs["mode"] = mode
  1602. if options:
  1603. self.kwargs["options"] = options
  1604. def __eq__(self, other):
  1605. return (isinstance(other, _TorchCompileWrapper) and
  1606. self.compiler_fn == other.compiler_fn and
  1607. self.kwargs == other.kwargs and
  1608. self.dynamic == other.dynamic)
  1609. def __call__(self, model_, inputs_):
  1610. return self.compiler_fn(model_, inputs_, **self.kwargs)
  1611. def reset(self):
  1612. if hasattr(self.compiler_fn, "reset"):
  1613. self.compiler_fn.reset()
  1614. def compile(model: Optional[Callable] = None, *,
  1615. fullgraph: builtins.bool = False,
  1616. dynamic: Optional[builtins.bool] = None,
  1617. backend: Union[str, Callable] = "inductor",
  1618. mode: Union[str, None] = None,
  1619. options: Optional[Dict[str, Union[str, builtins.int, builtins.bool]]] = None,
  1620. disable: builtins.bool = False) -> Callable:
  1621. """
  1622. Optimizes given model/function using TorchDynamo and specified backend.
  1623. If you are compiling an :class:`torch.nn.Module`, you can also use :meth:`torch.nn.Module.compile`
  1624. to compile the module inplace without changing its structure.
  1625. Concretely, for every frame executed within the compiled region, we will attempt
  1626. to compile it and cache the compiled result on the code object for future
  1627. use. A single frame may be compiled multiple times if previous compiled
  1628. results are not applicable for subsequent calls (this is called a "guard
  1629. failure), you can use TORCH_LOGS=guards to debug these situations.
  1630. Multiple compiled results can be associated with a frame up to
  1631. ``torch._dynamo.config.cache_size_limit``, which defaults to 8; at which
  1632. point we will fall back to eager. Note that compile caches are per
  1633. *code object*, not frame; if you dynamically create multiple copies of a
  1634. function, they will all share the same code cache.
  1635. Args:
  1636. model (Callable): Module/function to optimize
  1637. fullgraph (bool): If False (default), torch.compile attempts to discover compileable regions
  1638. in the function that it will optimize. If True, then we require that the entire function be
  1639. capturable into a single graph. If this is not possible (that is, if there are graph breaks),
  1640. then this will raise an error.
  1641. dynamic (bool or None): Use dynamic shape tracing. When this is True, we will up-front attempt
  1642. to generate a kernel that is as dynamic as possible to avoid recompilations when
  1643. sizes change. This may not always work as some operations/optimizations will
  1644. force specialization; use TORCH_LOGS=dynamic to debug overspecialization.
  1645. When this is False, we will NEVER generate dynamic kernels, we will always specialize.
  1646. By default (None), we automatically detect if dynamism has occurred and compile a more
  1647. dynamic kernel upon recompile.
  1648. backend (str or Callable): backend to be used
  1649. - "inductor" is the default backend, which is a good balance between performance and overhead
  1650. - Non experimental in-tree backends can be seen with `torch._dynamo.list_backends()`
  1651. - Experimental or debug in-tree backends can be seen with `torch._dynamo.list_backends(None)`
  1652. - To register an out-of-tree custom backend:
  1653. https://pytorch.org/docs/main/torch.compiler_custom_backends.html#registering-custom-backends
  1654. mode (str): Can be either "default", "reduce-overhead", "max-autotune" or "max-autotune-no-cudagraphs"
  1655. - "default" is the default mode, which is a good balance between performance and overhead
  1656. - "reduce-overhead" is a mode that reduces the overhead of python with CUDA graphs,
  1657. useful for small batches. Reduction of overhead can come at the cost of more memory
  1658. usage, as we will cache the workspace memory required for the invocation so that we
  1659. do not have to reallocate it on subsequent runs. Reduction of overhead is not guaranteed
  1660. to work; today, we only reduce overhead for CUDA only graphs which do not mutate inputs.
  1661. There are other circumstances where CUDA graphs are not applicable; use TORCH_LOG=perf_hints
  1662. to debug.
  1663. - "max-autotune" is a mode that leverages Triton based matrix multiplications and convolutions
  1664. It enables CUDA graphs by default.
  1665. - "max-autotune-no-cudagraphs" is a mode similar to "max-autotune" but without CUDA graphs
  1666. - To see the exact configs that each mode sets you can call `torch._inductor.list_mode_options()`
  1667. options (dict): A dictionary of options to pass to the backend. Some notable ones to try out are
  1668. - `epilogue_fusion` which fuses pointwise ops into templates. Requires `max_autotune` to also be set
  1669. - `max_autotune` which will profile to pick the best matmul configuration
  1670. - `fallback_random` which is useful when debugging accuracy issues
  1671. - `shape_padding` which pads matrix shapes to better align loads on GPUs especially for tensor cores
  1672. - `triton.cudagraphs` which will reduce the overhead of python with CUDA graphs
  1673. - `trace.enabled` which is the most useful debugging flag to turn on
  1674. - `trace.graph_diagram` which will show you a picture of your graph after fusion
  1675. - For inductor you can see the full list of configs that it supports by calling `torch._inductor.list_options()`
  1676. disable (bool): Turn torch.compile() into a no-op for testing
  1677. Example::
  1678. @torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
  1679. def foo(x):
  1680. return torch.sin(x) + torch.cos(x)
  1681. """
  1682. _C._log_api_usage_once("torch.compile")
  1683. if sys.version_info >= (3, 13):
  1684. raise RuntimeError("Dynamo is not supported on Python 3.13+")
  1685. # Decorator mode
  1686. if model is None:
  1687. def fn(model: Callable):
  1688. if model is None:
  1689. raise RuntimeError("Model can't be None")
  1690. return compile(model,
  1691. fullgraph=fullgraph,
  1692. dynamic=dynamic,
  1693. backend=backend,
  1694. mode=mode,
  1695. options=options,
  1696. disable=disable)
  1697. return fn
  1698. if mode is not None and options is not None:
  1699. raise RuntimeError("Either mode or options can be specified, but both can't be specified at the same time.")
  1700. if mode is None and options is None:
  1701. mode = "default"
  1702. if backend == "inductor":
  1703. backend = _TorchCompileInductorWrapper(mode, options, dynamic)
  1704. else:
  1705. backend = _TorchCompileWrapper(backend, mode, options, dynamic)
  1706. return torch._dynamo.optimize(backend=backend, nopython=fullgraph, dynamic=dynamic, disable=disable)(model)
  1707. from torch import export as export
  1708. from torch._higher_order_ops import cond
  1709. def _register_device_module(device_type, module):
  1710. r"""Register an external runtime module of the specific :attr:`device_type`
  1711. supported by torch.
  1712. After the :attr:`module` is registered correctly, the user can refer
  1713. the external runtime module as part of torch with attribute torch.xxx.
  1714. """
  1715. # Make sure the device_type represent a supported device type for torch.
  1716. device_type = torch.device(device_type).type
  1717. m = sys.modules[__name__]
  1718. if hasattr(m, device_type):
  1719. raise RuntimeError(f"The runtime module of '{device_type}' has already "
  1720. f"been registered with '{getattr(m, device_type)}'")
  1721. setattr(m, device_type, module)
  1722. torch_module_name = '.'.join([__name__, device_type])
  1723. sys.modules[torch_module_name] = module
  1724. # expose return_types
  1725. from . import return_types
  1726. from . import library
  1727. if not TYPE_CHECKING:
  1728. from . import _meta_registrations
  1729. # Enable CUDA Sanitizer
  1730. if 'TORCH_CUDA_SANITIZER' in os.environ:
  1731. import torch.cuda._sanitizer as csan
  1732. csan.enable_cuda_sanitizer()
  1733. # Populate magic methods on SymInt and SymFloat
  1734. import torch.fx.experimental.sym_node
  1735. from torch import func as func
  1736. from torch.func import vmap
  1737. # Register MPS specific decomps
  1738. torch.backends.mps._init()
  1739. if not _running_with_deploy():
  1740. from torch import compiler as compiler
  1741. class _TritonLibrary:
  1742. lib = torch.library.Library("triton", "DEF")
  1743. ops_table: Dict[Tuple[str, str], Callable] = {}
  1744. @classmethod
  1745. def registerOp(cls, op_key, full_schema, op_impl, dispatch_key):
  1746. if (op_key, dispatch_key) not in cls.ops_table:
  1747. cls.lib.define(full_schema)
  1748. cls.lib.impl("triton::" + op_key, op_impl, dispatch_key)
  1749. cls.ops_table[(op_key, dispatch_key)] = op_impl
  1750. return cls.ops_table[(op_key, dispatch_key)]
  1751. # Deprecated attributes
  1752. _deprecated_attrs = {
  1753. "has_mps": torch.backends.mps.is_built,
  1754. "has_cuda": torch.backends.cuda.is_built,
  1755. "has_cudnn": torch.backends.cudnn.is_available,
  1756. "has_mkldnn": torch.backends.mkldnn.is_available,
  1757. }
  1758. if TYPE_CHECKING:
  1759. # Import the following modules during type checking to enable code intelligence features,
  1760. # such as auto-completion in tools like pylance, even when these modules are not explicitly
  1761. # imported in user code.
  1762. from torch import _dynamo as _dynamo
  1763. from torch import _inductor as _inductor
  1764. from torch import onnx as onnx
  1765. else:
  1766. _lazy_modules = {
  1767. "_dynamo",
  1768. "_inductor",
  1769. "_export",
  1770. # ONNX must be imported after _dynamo, _ops, _subclasses, fx, func and jit
  1771. "onnx",
  1772. }
  1773. def __getattr__(name):
  1774. # Deprecated attrs
  1775. replacement = _deprecated_attrs.get(name)
  1776. if replacement is not None:
  1777. import warnings
  1778. warnings.warn(f"'{name}' is deprecated, please use '{replacement.__module__}.{replacement.__name__}()'", stacklevel=2)
  1779. return replacement()
  1780. # Lazy modules
  1781. if name in _lazy_modules:
  1782. import importlib
  1783. return importlib.import_module(f".{name}", __name__)
  1784. raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
  1785. def get_device_module(device: Optional[Union[torch.device, str]] = None):
  1786. """
  1787. Returns the module associated with a given device(e.g., torch.device('cuda'), "mtia:0", "xpu", ...).
  1788. If no device is given, return the module for the current accelerator or CPU if none is present.
  1789. """
  1790. if isinstance(device, torch.device):
  1791. device_module_name = device.type
  1792. elif isinstance(device, str):
  1793. device_module_name = torch.device(device).type
  1794. elif device is None:
  1795. # Using default accelerator type. If no accelerator is available, it automatically returns CPU device.
  1796. device_module_name = torch._C._get_accelerator().type
  1797. else:
  1798. raise RuntimeError(f"Invalid value of device '{device}', expect torch.device, str, or None")
  1799. device_module = getattr(torch, device_module_name, None)
  1800. if device_module is None:
  1801. raise RuntimeError(
  1802. f"Device '{device_module_name}' does not have a corresponding module registered as 'torch.{device_module_name}'."
  1803. )
  1804. return device_module
  1805. def _constrain_as_size(symbol, min: Optional[builtins.int] = None, max: Optional[builtins.int] = None):
  1806. """
  1807. This indicates that a given int is size-like, and can be used in any context where a size is expected.
  1808. You will typically use this when reading out integers from Tensors, e.g., max.item() or lengths.tolist()
  1809. which then need to be used as tensor constructors. Providing these assertions to PyTorch can help resolve
  1810. GuardOnDataDependentSymNode errors upon export, since we cannot guard on unbacked SymInts.
  1811. This function has unusual semantics in some circumstances in framework
  1812. code, we will treat this int as >= 2 (when we do a size-oblivious guard).
  1813. This makes it easier to use the unbacked int in size contexts,
  1814. as we will often attempt to guard on a size being zero/one
  1815. (e.g., when computing the contiguity of a tensor, or testing if
  1816. broadcasting can occur), which will not work on unbacked SymInts.
  1817. However, if we conservatively assume that the size is not zero/one, we will
  1818. end up with a graph that will still work even if the size is zero/one.
  1819. For more details, see https://docs.google.com/document/d/1HSuTTVvYH1pTew89Rtpeu84Ht3nQEFTYhAX3Ypa_xJs/edit
  1820. ```
  1821. """
  1822. torch.sym_constrain_range_for_size(symbol, min=min, max=max)
  1823. from . import _logging
  1824. _logging._init_logs()