builder.py 183 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340334133423343334433453346334733483349335033513352335333543355335633573358335933603361336233633364336533663367336833693370337133723373337433753376337733783379338033813382338333843385338633873388338933903391339233933394339533963397339833993400340134023403340434053406340734083409341034113412341334143415341634173418341934203421342234233424342534263427342834293430343134323433343434353436343734383439344034413442344334443445344634473448344934503451345234533454345534563457345834593460346134623463346434653466346734683469347034713472347334743475347634773478347934803481348234833484348534863487348834893490349134923493349434953496349734983499350035013502350335043505350635073508350935103511351235133514351535163517351835193520352135223523352435253526352735283529353035313532353335343535353635373538353935403541354235433544354535463547354835493550355135523553355435553556355735583559356035613562356335643565356635673568356935703571357235733574357535763577357835793580358135823583358435853586358735883589359035913592359335943595359635973598359936003601360236033604360536063607360836093610361136123613361436153616361736183619362036213622362336243625362636273628362936303631363236333634363536363637363836393640364136423643364436453646364736483649365036513652365336543655365636573658365936603661366236633664366536663667366836693670367136723673367436753676367736783679368036813682368336843685368636873688368936903691369236933694369536963697369836993700370137023703370437053706370737083709371037113712371337143715371637173718371937203721372237233724372537263727372837293730373137323733373437353736373737383739374037413742374337443745374637473748374937503751375237533754375537563757375837593760376137623763376437653766376737683769377037713772377337743775377637773778377937803781378237833784378537863787378837893790379137923793379437953796379737983799380038013802380338043805380638073808380938103811381238133814381538163817381838193820382138223823382438253826382738283829383038313832383338343835383638373838383938403841384238433844384538463847384838493850385138523853385438553856385738583859386038613862386338643865386638673868386938703871387238733874387538763877387838793880388138823883388438853886388738883889389038913892389338943895389638973898389939003901390239033904390539063907390839093910391139123913391439153916391739183919392039213922392339243925392639273928392939303931393239333934393539363937393839393940394139423943394439453946394739483949395039513952395339543955395639573958395939603961396239633964396539663967396839693970397139723973397439753976397739783979398039813982398339843985398639873988398939903991399239933994399539963997399839994000400140024003400440054006400740084009401040114012401340144015401640174018401940204021402240234024402540264027402840294030403140324033403440354036403740384039404040414042404340444045404640474048404940504051405240534054405540564057405840594060406140624063406440654066406740684069407040714072407340744075407640774078407940804081408240834084408540864087408840894090409140924093409440954096409740984099410041014102410341044105410641074108410941104111411241134114411541164117411841194120412141224123412441254126412741284129413041314132413341344135413641374138413941404141414241434144414541464147414841494150415141524153415441554156415741584159416041614162416341644165416641674168416941704171417241734174417541764177417841794180418141824183418441854186418741884189419041914192419341944195419641974198419942004201420242034204420542064207420842094210421142124213421442154216421742184219422042214222422342244225422642274228422942304231423242334234423542364237423842394240424142424243424442454246424742484249425042514252425342544255425642574258425942604261426242634264426542664267426842694270427142724273427442754276427742784279428042814282428342844285428642874288428942904291429242934294429542964297429842994300430143024303430443054306430743084309431043114312431343144315431643174318431943204321432243234324432543264327432843294330433143324333433443354336433743384339434043414342434343444345434643474348434943504351435243534354435543564357435843594360436143624363436443654366436743684369437043714372437343744375
  1. """
  2. This module contains classes and utilities for building variable trackers in Dynamo.
  3. Variable trackers are used to convert Python values into symbolic representations
  4. that can be traced and transformed during graph capture.
  5. The key classes are:
  6. - VariableBuilder: Handles source-tracked objects that need guards and proper
  7. reconstruction in the output graph. Used for inputs, module attributes, etc.
  8. - SourcelessBuilder: Handles ephemeral objects created during tracing that don't
  9. need source tracking or guards. Used for temporary lists, intermediate values, etc.
  10. Variable trackers enable Dynamo to track the flow of values through the program,
  11. maintain guards for dynamic properties, and reconstruct values in the output graph.
  12. The builders in this module handle converting Python values into appropriate
  13. VariableTracker instances based on their type and usage context.
  14. """
  15. import abc
  16. import collections
  17. import contextlib
  18. import copy
  19. import dataclasses
  20. import enum
  21. import functools
  22. import inspect
  23. import itertools
  24. import logging
  25. import math
  26. import operator
  27. import random
  28. import re
  29. import sys
  30. import time
  31. import types
  32. import weakref
  33. from collections.abc import Callable, MutableMapping
  34. from types import ModuleType
  35. from typing import Any, NamedTuple, NoReturn, Optional, overload, TYPE_CHECKING, Union
  36. import sympy
  37. import torch
  38. from torch import SymInt
  39. from torch._dispatch.python import enable_python_dispatcher
  40. from torch._dynamo.graph_bytecode_inputs import (
  41. get_external_object_by_index,
  42. register_user_object,
  43. )
  44. from torch._dynamo.utils import (
  45. get_metrics_context,
  46. is_int_specialization_case,
  47. is_torch_sym,
  48. set_feature_use,
  49. )
  50. from torch._guards import TracingContext
  51. from torch._higher_order_ops.flat_apply import flat_apply
  52. from torch._higher_order_ops.torchbind import call_torchbind
  53. from torch._library.opaque_object import (
  54. is_opaque_reference_type,
  55. is_opaque_type,
  56. is_opaque_value_type,
  57. should_hoist,
  58. )
  59. from torch._ops import HigherOrderOperator, OpOverload, OpOverloadPacket
  60. from torch._subclasses.fake_tensor import (
  61. FakeTensor,
  62. FakeTensorMode,
  63. is_fake,
  64. maybe_get_fake_mode,
  65. )
  66. from torch._subclasses.meta_utils import is_sparse_any, safe_grad
  67. from torch._utils_internal import justknobs_check
  68. from torch.fx.experimental._backward_state import BackwardState
  69. from torch.fx.experimental._dynamism import normalize_source_name
  70. from torch.fx.experimental.sym_node import _DynamicScalar, DynamicInt
  71. from torch.fx.experimental.symbolic_shapes import (
  72. _constrain_range_for_size,
  73. _nested_int_aware_sort,
  74. DimDynamic,
  75. RelaxedUnspecConstraint,
  76. StatefulSymbolicContext,
  77. SubclassSymbolicContext,
  78. SymbolicContext,
  79. SymIntSymbolicContext,
  80. TrackedFake,
  81. )
  82. from torch.fx.immutable_collections import immutable_dict, immutable_list
  83. from torch.nn.utils._expanded_weights import ExpandedWeight
  84. from torch.utils._ordered_set import OrderedSet
  85. from torch.utils._python_dispatch import (
  86. is_traceable_wrapper_subclass,
  87. is_traceable_wrapper_subclass_type,
  88. )
  89. from torch.utils._sympy.value_ranges import ValueRanges
  90. from torch.utils.weak import TensorWeakRef
  91. from .. import config, graph_break_hints, mutation_guard, replay_record, trace_rules
  92. from ..device_interface import get_registered_device_interfaces
  93. from ..exc import InternalTorchDynamoError, raise_observed_exception, unimplemented
  94. from ..guards import GuardBuilder, install_guard, make_dupe_guard
  95. from ..pgo import (
  96. auto_dynamic,
  97. auto_unset,
  98. FrameStateSizeEntry,
  99. InferStride,
  100. process_automatic_dynamic,
  101. )
  102. from ..side_effects import SideEffects
  103. from ..source import (
  104. AttrProxySource,
  105. AttrSource,
  106. CallMethodItemSource,
  107. ChainedSource,
  108. ConstDictKeySource,
  109. ConvertIntSource,
  110. DictGetItemSource,
  111. DictSubclassGetItemSource,
  112. DynamicScalarSource,
  113. FloatTensorSource,
  114. GetItemSource,
  115. GradSource,
  116. is_constant_source,
  117. is_from_closure_source,
  118. is_from_global_source,
  119. is_from_nonlocal_source,
  120. is_from_optimizer_source,
  121. is_from_unspecialized_nn_module_source,
  122. ListGetItemSource,
  123. LocalSource,
  124. NonSerializableSetGetItemSource,
  125. NumpyTensorSource,
  126. OptimizerSource,
  127. RandomValueSource,
  128. SkipGuardSource,
  129. Source,
  130. SubclassAttrListSource,
  131. TupleIteratorGetItemSource,
  132. UnspecializedBuiltinNNModuleSource,
  133. UnspecializedNNModuleSource,
  134. )
  135. from ..utils import (
  136. _extract_tensor_dict,
  137. build_checkpoint_variable,
  138. build_invoke_subgraph_variable,
  139. clone_input,
  140. common_constant_types,
  141. dict_keys,
  142. get_fake_value,
  143. get_items_from_dict,
  144. get_locals_to_steal,
  145. get_static_address_type,
  146. is_frozen_dataclass,
  147. is_function,
  148. is_function_or_wrapper,
  149. is_invoke_subgraph,
  150. is_lru_cache_wrapped_function,
  151. is_namedtuple,
  152. is_parameter_freezing,
  153. is_typing,
  154. is_utils_checkpoint,
  155. is_wrapper_or_member_descriptor,
  156. istype,
  157. namedtuple_fields,
  158. odict_values,
  159. proxy_args_kwargs,
  160. range_iterator,
  161. set_example_value,
  162. tensor_always_has_static_shape,
  163. tuple_iterator,
  164. tuple_iterator_getitem,
  165. tuple_iterator_len,
  166. unwrap_with_attr_name_if_wrapper,
  167. wrap_fake_exception,
  168. )
  169. from .base import (
  170. AttributeMutationNew,
  171. typestr,
  172. ValueMutationExisting,
  173. ValueMutationNew,
  174. VariableTracker,
  175. VariableTrackerMeta,
  176. )
  177. from .builtin import BuiltinVariable
  178. from .constant import ConstantVariable, EnumVariable
  179. from .ctx_manager import (
  180. AutocastModeVariable,
  181. DynamoConfigPatchVariable,
  182. ErrorOnGraphBreakVariable,
  183. NullContextVariable,
  184. PreserveVersionContextVariable,
  185. )
  186. from .dicts import (
  187. ConstDictVariable,
  188. DefaultDictVariable,
  189. DictKeySetVariable,
  190. FrozensetVariable,
  191. MappingProxyVariable,
  192. OrderedSetClassVariable,
  193. OrderedSetVariable,
  194. SetVariable,
  195. )
  196. from .distributed import (
  197. DeviceMeshVariable,
  198. PlacementClassVariable,
  199. PlacementVariable,
  200. ProcessGroupVariable,
  201. WorldMetaClassVariable,
  202. )
  203. from .functions import (
  204. BuiltinMethodVariable,
  205. CollectionsNamedTupleFunction,
  206. CollectiveFunctionRewriteVariable,
  207. CreateTMADescriptorExperimentalVariable,
  208. CreateTMADescriptorStableVariable,
  209. FunctoolsPartialVariable,
  210. FunctoolsWrapsVariable,
  211. SysFunctionVariable,
  212. TritonKernelVariable,
  213. TritonSetAllocatorSkipVariable,
  214. UserFunctionVariable,
  215. UserMethodVariable,
  216. WrapperUserFunctionVariable,
  217. )
  218. from .higher_order_ops import (
  219. LocalMapWrappedHigherOrderVariable,
  220. TorchHigherOrderOperatorVariable,
  221. )
  222. from .iter import ItertoolsVariable
  223. from .lazy import LazyConstantVariable, LazyVariableTracker
  224. from .lists import (
  225. BaseListVariable,
  226. ListIteratorVariable,
  227. ListVariable,
  228. NamedTupleVariable,
  229. RangeVariable,
  230. SizeVariable,
  231. SliceVariable,
  232. TupleIteratorVariable,
  233. TupleVariable,
  234. )
  235. from .misc import (
  236. AutogradEngineVariable,
  237. AutogradFunctionContextVariable,
  238. AutogradFunctionVariable,
  239. ComptimeVariable,
  240. ConstantLikeVariable,
  241. DebuggingVariable,
  242. DelayGraphBreakVariable,
  243. GetAttrVariable,
  244. GetSetDescriptorVariable,
  245. IgnoredFunctionVariable,
  246. LambdaVariable,
  247. LoggingLoggerVariable,
  248. MethodWrapperVariable,
  249. NumpyDTypeVariable,
  250. NumpyVariable,
  251. ObjectVariable,
  252. PythonModuleVariable,
  253. RandomClassVariable,
  254. RandomVariable,
  255. SavedTensorBox,
  256. TorchVersionVariable,
  257. TypingVariable,
  258. WeakRefVariable,
  259. )
  260. from .nn_module import (
  261. FSDPManagedNNModuleVariable,
  262. UnspecializedBuiltinNNModuleVariable,
  263. UnspecializedNNModuleVariable,
  264. )
  265. from .optimizer import OptimizerVariable
  266. from .script_object import OpaqueObjectClassVariable, TorchScriptObjectVariable
  267. from .sdpa import SDPAParamsVariable
  268. from .streams import EventVariable, StreamContextVariable, StreamVariable
  269. from .tensor import (
  270. NumpyNdarrayVariable,
  271. supported_const_comparison_op_values,
  272. SymNodeVariable,
  273. TensorSubclassVariable,
  274. TensorVariable,
  275. UnspecializedPythonVariable,
  276. )
  277. from .torch import (
  278. DispatchKeySetVariable,
  279. FuncTorchInterpreterVariable,
  280. TorchCtxManagerClassVariable,
  281. TorchInGraphFunctionVariable,
  282. )
  283. from .torch_function import (
  284. TensorWithTFOverrideVariable,
  285. torch_function_mode_stack_state_mgr,
  286. TorchFunctionModeVariable,
  287. )
  288. from .user_defined import (
  289. FrozenDataClassVariable,
  290. InspectVariable,
  291. IntWrapperVariable,
  292. KeyedJaggedTensorVariable,
  293. MutableMappingVariable,
  294. SourcelessGraphModuleVariable,
  295. UserDefinedClassVariable,
  296. UserDefinedDictVariable,
  297. UserDefinedEnumClassVariable,
  298. UserDefinedExceptionClassVariable,
  299. UserDefinedListVariable,
  300. UserDefinedObjectVariable,
  301. UserDefinedSetVariable,
  302. UserDefinedTupleVariable,
  303. )
  304. try:
  305. import numpy as np
  306. except ModuleNotFoundError:
  307. np: ModuleType = None # type: ignore[assignment]
  308. if TYPE_CHECKING:
  309. from torch._dynamo.codegen import PyCodegen
  310. from torch._dynamo.symbolic_convert import (
  311. InstructionTranslator,
  312. InstructionTranslatorBase,
  313. )
  314. log = logging.getLogger(__name__)
  315. static_inputs_log = torch._logging.getArtifactLogger(
  316. __name__, "cudagraph_static_inputs"
  317. )
  318. from typing import TypeVar
  319. # Placeholder for a VariableTracker to be used in proxy
  320. # creation so that we don't type erase
  321. VTTypeAlias = TypeVar("VTTypeAlias")
  322. T = TypeVar("T")
  323. DimList = list
  324. def safe_has_grad(t: object) -> bool:
  325. with torch._logging.hide_warnings(torch._logging._internal.safe_grad_filter):
  326. return hasattr(t, "grad")
  327. class _missing:
  328. pass
  329. @dataclasses.dataclass
  330. class GraphArg:
  331. source: Source | None
  332. # TODO: storing a SymInt here but not a FakeTensor is a pretty strange
  333. # thing to do. Probably should have example (which stores an int) and
  334. # fake_example
  335. _example: Any
  336. # When True, this indicates that this GraphArg is a Python quantity (e.g.,
  337. # a float or int) which we pass to the FX graph as a Tensor. This
  338. # controls how we codegen calls into the Dynamo graph: we will call
  339. # torch.as_tensor on the quantity before passing it in.
  340. #
  341. # Note that we typically do not pass dynamic integers as tensors, because
  342. # they will most frequently just be used for size computation. But this
  343. # is a policy decision that we can change our mind on; in particular, when
  344. # an int comes from a random number generator (e.g., random.randint), we
  345. # DO pass it as a tensor.
  346. #
  347. # It's also worth noting that our current tracing rules for
  348. # pass_arg_as_tensor as subtly broken: we just pun the variable as a
  349. # 0d scalar Tensor and pray that the semantics are the same. Which they
  350. # often are, but not necessarily. ezyang(May 2024) plans to fix this
  351. # soon.
  352. pass_arg_as_tensor: bool
  353. fake_tensor: torch._subclasses.fake_tensor.FakeTensor | None
  354. # UnspecializedPythonVariable often masquerades as a tensor.
  355. # We MUST NOT generate shape guard code
  356. # that actually tries to access tensor properties on these values.
  357. # is_tensor lets us tell if this graph arg actually is a tensor
  358. # or not.
  359. is_tensor: bool = True
  360. # Sometimes, the Tensor we pass to example is freshly allocated (smh).
  361. # Then we cannot only keep a weak reference to it. This lets you
  362. # stash a strong reference too.
  363. example_strong_ref: torch.Tensor | torch.SymInt | None = None
  364. def __setattr__(self, name: str, value: Any) -> None:
  365. # Use object.__setattr__ to bypass Dynamo's STORE_ATTR interception.
  366. # This is needed because when PYTORCH_TEST_WITH_DYNAMO=1, even internal
  367. # GraphArg creation can be traced, and with replay_side_effects=False,
  368. # normal STORE_ATTR bytecode only records mutations without applying them.
  369. object.__setattr__(self, name, value)
  370. @property
  371. def example(self) -> torch.Tensor | torch.SymInt | BackwardState | None:
  372. if isinstance(self._example, TensorWeakRef):
  373. r = self._example()
  374. assert r is not None
  375. return r
  376. else:
  377. return self._example
  378. def __post_init__(self) -> None:
  379. if isinstance(self._example, torch.Tensor):
  380. self._example = TensorWeakRef(self._example)
  381. assert is_fake(self.fake_tensor)
  382. def reconstruct(self, codegen: "PyCodegen") -> None:
  383. codegen(self.source)
  384. def erase(self) -> None:
  385. self._example = None
  386. self.example_strong_ref = None
  387. def __eq__(self, other: object) -> bool:
  388. if not isinstance(other, GraphArg):
  389. return False
  390. if self.source is None:
  391. return other.source is None
  392. else:
  393. if other.source is None:
  394. return False
  395. return self.source.name == other.source.name
  396. class BackwardStateGraphArg(GraphArg):
  397. def __init__(self) -> None:
  398. super().__init__(
  399. source=None,
  400. _example=BackwardState(),
  401. pass_arg_as_tensor=False,
  402. fake_tensor=None,
  403. is_tensor=False,
  404. )
  405. def reconstruct(self, codegen: "PyCodegen") -> None:
  406. assert codegen.tx.output.backward_state_var
  407. codegen.add_push_null(
  408. lambda: codegen.load_import_from(BackwardState.__module__, "BackwardState")
  409. )
  410. codegen.call_function(0, False)
  411. codegen.dup_top()
  412. codegen.store(codegen.tx.output.backward_state_var)
  413. # All class-based iterators in itertools
  414. # NOTE: use id() because some objects are not hashable, it will raise error during lookup
  415. ITERTOOLS_TYPE_IDS: frozenset[int] = frozenset(
  416. id(member)
  417. for name, member in vars(itertools).items()
  418. if not name.startswith("_") and inspect.isclass(member)
  419. )
  420. # Will be updated later in substitute_in_graph in torch/_dynamo/polyfills/itertools.py
  421. ITERTOOLS_POLYFILLED_TYPE_IDS: set[int] = set()
  422. # Capture fn pointer at import time
  423. # This is to guard against trying to mark the iterated tensors
  424. # as static in case user overrides fn ptr
  425. og_module_named_buffers_fn_ptr = torch.nn.Module.named_buffers
  426. og_module_named_parameters_fn_ptr = torch.nn.Module.named_parameters
  427. class VariableBuilder:
  428. """Wrap a python value in a VariableTracker() instance"""
  429. def __init__(
  430. self,
  431. tx: "InstructionTranslator",
  432. source: Source,
  433. allow_lazy_constant: bool = True,
  434. ) -> None:
  435. assert source is not None, (
  436. "Consider SourcelessBuilder for ephemeral objects, usually objects created locally."
  437. )
  438. assert TracingContext.try_get() is not None, "Expected active TracingContext"
  439. super().__init__()
  440. self.tx = tx
  441. self.source = source
  442. self.name = source.name
  443. # allow_lazy_constant controls whether LazyConstantVariable can be returned
  444. # for int/float/bool/str. Set to False when called from LazyCache.realize()
  445. # to prevent double-wrapping (LazyVariableTracker containing LazyConstantVariable).
  446. self.allow_lazy_constant = allow_lazy_constant
  447. def __call__(self, value: object) -> VariableTracker:
  448. _t0 = time.time_ns()
  449. try:
  450. return self._call_impl(value)
  451. finally:
  452. self.tx.output.bytecode_tracing_timings.variable_builder_call_ns += (
  453. time.time_ns() - _t0
  454. )
  455. def _call_impl(self, value: object) -> VariableTracker:
  456. if value in self.tx.output.side_effects:
  457. side_effect_result = self.tx.output.side_effects[value]
  458. dup_guard = make_dupe_guard(self.source, side_effect_result.source)
  459. if dup_guard:
  460. self.install_guards(dup_guard)
  461. if isinstance(value, torch.nn.Module) and isinstance(
  462. side_effect_result, UnspecializedNNModuleVariable
  463. ):
  464. # This means that two nn module instances with different sources
  465. # have the same id. NN modules are somewhat special objects,
  466. # because we have to track their nn_module_stack for ease of
  467. # use. But if we don't do anything, we will just return the
  468. # older variable tracker with the older nn_module_stack. So,
  469. # lets return the old variable tracker but update its
  470. # nn_module_stack
  471. side_effect_result.set_nn_module_stack_source(self.source)
  472. return side_effect_result
  473. cached_vt = self.tx.output.variable_tracker_cache.get(self.source)
  474. if cached_vt:
  475. # If allow_lazy_constant=False but the cached VT is a lazy variable,
  476. # we need to rebuild to get a non-lazy version. This happens when
  477. # LazyConstantVariable.realize() calls VariableBuilder.
  478. if self.allow_lazy_constant or not isinstance(
  479. cached_vt, LazyVariableTracker
  480. ):
  481. return cached_vt
  482. vt = self._wrap(value)
  483. if vt.source is None:
  484. vt.source = self.source
  485. def _is_deduplicable_sym_variable(value: Any, vt: VariableTracker) -> bool:
  486. # Constants like 0, 1, 2, etc. can be unspecialized as SymNodeVariables sometimes, but we
  487. # should NOT track them. If we use a single SymNodeVariable instance to track them
  488. # across multiple uses, then guards created for one usage will incorrectly apply to
  489. # all other usages of that constant, leading to unnecessary recompilations.
  490. return (
  491. is_torch_sym(value) or isinstance(value, _DynamicScalar)
  492. ) and isinstance(vt, SymNodeVariable)
  493. if (
  494. (
  495. self._can_lift_attrs_to_inputs(vt)
  496. or _is_deduplicable_sym_variable(value, vt)
  497. )
  498. and value not in self.tx.output.side_effects
  499. and not is_wrapper_or_member_descriptor(value)
  500. ):
  501. vt = self.tx.output.side_effects.track_object_existing(value, vt)
  502. # Skip caching for JVP_NESTING source because
  503. # JvpIncrementNestingCtxManagerVariable hides global JVP mutation from
  504. # Dynamo, resulting in stale value. We attempted a fix in
  505. # https://github.com/pytorch/pytorch/pull/174329 but it exposed other
  506. # issues. This only affects cache hit rate, NOT correctness.
  507. if "JVP_NESTING" not in self.source.name:
  508. self.tx.output.variable_tracker_cache[self.source] = vt
  509. return vt
  510. def _can_lift_attrs_to_inputs(self, vt: VariableTracker) -> bool:
  511. return type(vt) in {
  512. TensorVariable,
  513. TensorWithTFOverrideVariable,
  514. UserDefinedObjectVariable,
  515. NumpyNdarrayVariable,
  516. TorchScriptObjectVariable,
  517. }
  518. def get_source(self) -> Source:
  519. return self.source
  520. def install_guards(self, *guards: Callable[..., Any]) -> dict[str, Any] | None:
  521. source = self.get_source()
  522. try:
  523. tmp = [source.make_guard(guard) for guard in guards]
  524. except NotImplementedError:
  525. return None
  526. install_guard(*tmp, skip=1)
  527. return {}
  528. @classmethod
  529. def _type_dispatch(cls) -> dict[object, Callable[..., Any]]:
  530. return cls._type_dispatch_impl(config.trace_numpy)
  531. @classmethod
  532. @functools.cache
  533. def _type_dispatch_impl(cls, trace_numpy: bool) -> dict[object, Callable[..., Any]]:
  534. # NB: Careful not to close over self to avoid ref cycle from lru_cache
  535. entries = [
  536. (
  537. (
  538. torch.Tensor,
  539. torch.nn.Parameter,
  540. torch._subclasses.FakeTensor,
  541. torch._subclasses.functional_tensor.FunctionalTensor,
  542. ),
  543. cls.wrap_tensor,
  544. ),
  545. (
  546. (tuple, list, odict_values, collections.deque, torch.Size),
  547. cls.wrap_listlike,
  548. ),
  549. (tuple_iterator, cls.wrap_tuple_iterator),
  550. (range_iterator, cls.wrap_range_iterator),
  551. ((slice, range), cls.wrap_slice_range),
  552. (tuple(common_constant_types), cls.wrap_literal),
  553. (re.Pattern, cls.wrap_regex_pattern),
  554. (weakref.ReferenceType, cls.wrap_weakref),
  555. (torch.utils.hooks.RemovableHandle, cls.wrap_removable_handle),
  556. (torch.jit.ScriptFunction, cls.wrap_jit_function),
  557. (types.MappingProxyType, cls.wrap_mapping_proxy),
  558. ]
  559. if trace_numpy and np:
  560. # pyrefly: ignore [bad-argument-type]
  561. entries.append((np.ndarray, cls.wrap_numpy_ndarray))
  562. # pyrefly: ignore [implicit-any]
  563. result = {}
  564. for ts, fn in entries:
  565. for t in ts if isinstance(ts, tuple) else (ts,):
  566. assert t not in result
  567. result[t] = fn
  568. return result
  569. def wrap_regex_pattern(self, value: re.Pattern[Any]) -> ConstantLikeVariable:
  570. # TODO(jansel): something like a REPR_MATCH might be more robust here
  571. self.install_guards(GuardBuilder.ID_MATCH)
  572. return ConstantLikeVariable(value)
  573. def wrap_weakref(self, value: weakref.ReferenceType[Any]) -> WeakRefVariable:
  574. self.install_guards(GuardBuilder.TYPE_MATCH)
  575. return WeakRefVariable.build(self.tx, value, source=self.source)
  576. def wrap_removable_handle(
  577. self, value: torch.utils.hooks.RemovableHandle
  578. ) -> NoReturn:
  579. # This means that the removable handle was created in some other frame.
  580. # Our current infra requires the hook to be registered and removed in
  581. # the same frame. So graph break.
  582. # Related test - PYTORCH_TEST_WITH_DYNAMO=1 python test/test_autograd.py -k TestAutograd.test_hooks
  583. unimplemented(
  584. gb_type="Attempted to represent unregistered RemovableHandle",
  585. context="",
  586. explanation="Dynamo attempted to build a representation of a torch.utils.hooks.RemovableHandle, "
  587. "which is not supported. This happens because the RemovableHandle was created in another frame.",
  588. hints=[],
  589. )
  590. def wrap_jit_function(self, value: Any) -> WrapperUserFunctionVariable:
  591. self.install_guards(GuardBuilder.TYPE_MATCH)
  592. return WrapperUserFunctionVariable(
  593. value, "_torchdynamo_inline", source=self.source
  594. )
  595. def wrap_mapping_proxy(self, value: Any) -> VariableTracker:
  596. self.install_guards(GuardBuilder.TYPE_MATCH)
  597. # This might be suboptimal compared to dict guards. But mappingproxy is
  598. # not very common, so its ok to guard on all keys.
  599. self.install_guards(GuardBuilder.MAPPING_KEYS_CHECK)
  600. all_const = all(ConstantVariable.is_literal(k) for k in value)
  601. if not all_const:
  602. unimplemented(
  603. gb_type="non-const keys in mappingproxy",
  604. context=f"non-const keys: {[k for k in value.keys() if not ConstantVariable.is_literal(k)]}", # noqa: SIM118
  605. explanation="Dynamo expects mappingproxy keys to be constants.",
  606. hints=[
  607. "Ensure your mappingproxy keys are constants (e.g. int, float, strings)",
  608. ],
  609. )
  610. def build_key_value(k: Any, v: Any) -> tuple[VariableTracker, VariableTracker]:
  611. key = ConstantVariable.create(k)
  612. source_key = k
  613. source_value = GetItemSource(self.get_source(), source_key)
  614. res_value = LazyVariableTracker.create(v, source_value)
  615. return key, res_value
  616. items = dict(build_key_value(k, v) for k, v in value.items())
  617. # Create a dict_vt to be used in the mapping proxy variable
  618. # pyrefly: ignore[bad-argument-type]
  619. dict_vt = ConstDictVariable(items, source=None)
  620. result = MappingProxyVariable(dict_vt, source=self.source)
  621. return self.tx.output.side_effects.track_mutable(value, result)
  622. @classmethod
  623. @functools.cache
  624. def _id_dispatch(
  625. cls,
  626. ) -> dict[int, Callable[["VariableBuilder", Any], VariableTracker]]:
  627. from ..comptime import comptime
  628. entries = [
  629. (comptime, lambda self, value: ComptimeVariable()),
  630. (
  631. dataclasses.fields,
  632. lambda self, value: LambdaVariable(
  633. _dataclasses_fields_lambda,
  634. source=self.source,
  635. **self.install_guards(GuardBuilder.CLOSURE_MATCH),
  636. ),
  637. ),
  638. (torch.__version__, lambda self, value: TorchVersionVariable()),
  639. ]
  640. # pyrefly: ignore [implicit-any]
  641. result = {}
  642. for ts, fn in entries:
  643. for t in ts if isinstance(ts, (tuple, list)) else (ts,):
  644. assert t not in result
  645. result[id(t)] = fn
  646. return result
  647. def _wrap(self, value: Any) -> VariableTracker:
  648. # import here to avoid circular dependencies
  649. from torch.utils._triton import (
  650. has_triton,
  651. has_triton_experimental_host_tma,
  652. has_triton_tensor_descriptor_host_tma,
  653. )
  654. from ..decorators import (
  655. DynamoConfigPatchProxy,
  656. ErrorOnGraphBreakDecoratorContextManager,
  657. )
  658. if has_triton():
  659. from triton.runtime.autotuner import Autotuner
  660. from triton.runtime.jit import JITFunction
  661. else:
  662. class JITFunction:
  663. pass
  664. class Autotuner:
  665. pass
  666. # default implementations, in case we don't have triton (or the wrong triton version)
  667. def create_1d_tma_descriptor() -> None:
  668. pass
  669. def create_2d_tma_descriptor() -> None:
  670. pass
  671. class TensorDescriptor:
  672. @staticmethod
  673. def from_tensor() -> None:
  674. pass
  675. def set_allocator() -> None:
  676. pass
  677. if has_triton_experimental_host_tma():
  678. from triton.tools.experimental_descriptor import ( # noqa: F811
  679. create_1d_tma_descriptor,
  680. create_2d_tma_descriptor,
  681. )
  682. if has_triton_tensor_descriptor_host_tma():
  683. from triton.tools.tensor_descriptor import TensorDescriptor # noqa: F811
  684. if has_triton():
  685. import triton as triton_mod
  686. if hasattr(triton_mod, "set_allocator"):
  687. set_allocator = triton_mod.set_allocator # noqa: F811
  688. # Handle exact type() match
  689. type_dispatch = self._type_dispatch().get(type(value))
  690. if type_dispatch is not None:
  691. return type_dispatch(self, value)
  692. # Handle exact id() match
  693. id_dispatch = self._id_dispatch().get(id(value))
  694. if id_dispatch is not None:
  695. return id_dispatch(self, value)
  696. # Everything else (NB: order matters!)
  697. if (
  698. isinstance(value, torch.Tensor)
  699. and type(value)
  700. not in (
  701. # These torch-native subclasses have overly restrictive
  702. # `__torch_function__` which prevents Dynamo from reading their
  703. # tensor attributes like `is_nested` or calling methods like
  704. # `_is_view`.
  705. torch.nn.parameter.UninitializedBuffer,
  706. torch.nn.parameter.UninitializedParameter,
  707. ExpandedWeight,
  708. )
  709. and type(value) not in config.nontraceable_tensor_subclasses
  710. ):
  711. if (
  712. type(value).__torch_dispatch__ is torch.Tensor.__torch_dispatch__
  713. or is_traceable_wrapper_subclass(value)
  714. ):
  715. return self.wrap_tensor(value)
  716. if is_namedtuple(value):
  717. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  718. output: list[VariableTracker] = [
  719. LazyVariableTracker.create(
  720. getattr(value, name),
  721. source=AttrSource(self.source, name),
  722. )
  723. for name in namedtuple_fields(type(value))
  724. ]
  725. tuple_vt = TupleVariable(
  726. output,
  727. source=self.source,
  728. mutation_type=ValueMutationExisting(),
  729. )
  730. result = NamedTupleVariable(
  731. output,
  732. tuple_cls=type(value),
  733. source=self.source,
  734. tuple_vt=tuple_vt,
  735. )
  736. return self.tx.output.side_effects.track_object_existing(value, result)
  737. elif istype(value, (dict, collections.defaultdict, collections.OrderedDict)):
  738. self.install_guards(GuardBuilder.TYPE_MATCH)
  739. all_const = all(ConstantVariable.is_literal(k) for k in value)
  740. # For all_const, we don't have to guard on anything yet. We guard on
  741. # keys lazily by adding a dict_getitem entry for each accessed key.
  742. # For cases where we need to guard on all keys, we lazily put guards
  743. # during the dict call_method (check dicts.py)
  744. if not all_const:
  745. # Guard on the key order
  746. # This is not ideal, i.e., there is no need to guard on the key
  747. # order. But we guard on the key order because of the complexity
  748. #
  749. # 1) For non-constant objects, we can't save the key in the
  750. # guard context because it can be memory heavy. We can add
  751. # weakrefs but this complicates the accesses.
  752. #
  753. # 2) For non-constant objects, we also have to guard on the keys
  754. # (like TENSOR_MATCH on tensor). We might also have guards on
  755. # the attributes of the keys (like tensor.grad). To make this
  756. # work in tree structure is complicated.
  757. #
  758. # So, instead we guard on the key order. While guarding on key
  759. # order, we just save the indices and use it to access keys and
  760. # values. Indices are cheap to save.
  761. self.tx.output.guard_on_key_order.add(self.source)
  762. # We need all the keys to be hashable. We do this within the
  763. # _HashableTracker class in dicts.py
  764. def build_key_value(
  765. i: Any, k: Any, v: Any
  766. ) -> tuple[VariableTracker, VariableTracker]:
  767. base = self.get_source()
  768. if all_const:
  769. key = ConstantVariable.create(k)
  770. source_key = k
  771. else:
  772. source_key = ConstDictKeySource(base, i)
  773. key = LazyVariableTracker.create(k, source_key)
  774. source_value = DictGetItemSource(base, source_key)
  775. res_value = LazyVariableTracker.create(v, source_value)
  776. return key, res_value
  777. # Ensure that we call dict.keys and not value.keys (which can call
  778. # overridden keys method). In the C++ guards, we relied on
  779. # PyDict_Next to traverse the dictionary, which uses the internal
  780. # data structure and does not call the overridden keys method.
  781. result = dict(
  782. build_key_value(i, k, v)
  783. for i, (k, v) in enumerate(get_items_from_dict(value))
  784. )
  785. if istype(value, collections.defaultdict):
  786. factory_source = AttrSource(self.source, "default_factory")
  787. result = DefaultDictVariable(
  788. result, # type: ignore[arg-type]
  789. type(value),
  790. default_factory=VariableBuilder(self.tx, factory_source)(
  791. value.default_factory
  792. ),
  793. source=self.source,
  794. )
  795. else:
  796. result = ConstDictVariable(
  797. result, # type: ignore[arg-type]
  798. user_cls=type(value),
  799. source=self.source,
  800. )
  801. return self.tx.output.side_effects.track_mutable(value, result)
  802. elif isinstance(value, torch.nn.Module):
  803. return self.wrap_module(value)
  804. elif ConstantVariable.is_literal(value): # non-atomic literals
  805. return self.wrap_literal(value)
  806. elif isinstance(value, torch.overrides.TorchFunctionMode):
  807. var = TorchFunctionModeVariable(value, source=self.source)
  808. self.tx.output.side_effects.track_object_existing(value, var)
  809. return var
  810. elif istype(value, (set, OrderedSet)):
  811. if any(isinstance(x, torch.Tensor) for x in value):
  812. unimplemented(
  813. gb_type="Attempted to wrap a set with tensors",
  814. context="Python set containing torch.Tensor elements",
  815. explanation=(
  816. "Dynamo cannot trace sets of tensors. To get a stable ordering, "
  817. "Dynamo needs to convert the set into a list and the order might not be "
  818. "stable if the set contains tensors."
  819. ),
  820. hints=[
  821. "Use a dictionary where the keys are tensors.",
  822. *graph_break_hints.SUPPORTABLE,
  823. ],
  824. )
  825. self.install_guards(GuardBuilder.TYPE_MATCH)
  826. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  827. set_var_cls = SetVariable
  828. if istype(value, OrderedSet):
  829. # Guard on the internal dict of OrderedSet
  830. internal_dict_source = AttrSource(self.source, "_dict")
  831. install_guard(
  832. internal_dict_source.make_guard(GuardBuilder.DICT_KEYS_MATCH)
  833. )
  834. self.tx.output.guard_on_key_order.add(internal_dict_source)
  835. set_var_cls = OrderedSetVariable
  836. # The list gives a ordering for the set items. The ordering is based
  837. # on the Python hash and it is not related to object ordering inside
  838. # the set object. The order being incorrect at runtime will lead to
  839. # a recompilation.
  840. L = list(value)
  841. items = [
  842. LazyVariableTracker.create(
  843. v, source=NonSerializableSetGetItemSource(self.source, i)
  844. )
  845. for i, v in enumerate(L)
  846. ]
  847. result = set_var_cls(items, source=self.source)
  848. return self.tx.output.side_effects.track_object_existing(value, result)
  849. elif istype(value, frozenset) and all(
  850. (
  851. # For DBR quantization, we could get a frozenset of torch funcs.
  852. (type(x) is types.BuiltinMethodType and x.__module__ == "torch")
  853. or
  854. # Another commonly used frozenset of types.
  855. x in torch.utils._pytree.BUILTIN_TYPES
  856. or
  857. # For activation checkpointing, we could get a frozenset of torch ops.
  858. isinstance(x, (OpOverload, OpOverloadPacket))
  859. )
  860. for x in value
  861. ):
  862. # For the limited cases of frozenset here, we know the items won't
  863. # change across runs, so we can safely create sourceless VTs for
  864. # them and guard on the frozenset contents via EQUALS_MATCH.
  865. # TODO support source for sets and remove the special logics here.
  866. items = [SourcelessBuilder.create(self.tx, v) for v in value]
  867. self.install_guards(GuardBuilder.EQUALS_MATCH)
  868. return FrozensetVariable(items, source=self.source)
  869. elif isinstance(
  870. value, (enum.Enum, torch.DispatchKey, torch._C._functorch.TransformType)
  871. ):
  872. self.install_guards(GuardBuilder.ID_MATCH)
  873. return EnumVariable(value=value, source=self.source)
  874. elif DebuggingVariable.is_reorderable_logging_function(value):
  875. # Put this above builtin_callable so that print() can be handled
  876. # along with other builtin debugging functions
  877. self.install_guards(GuardBuilder.BUILTIN_MATCH)
  878. return DebuggingVariable(value, source=self.source)
  879. elif callable(value) and any(
  880. value is fn for fn in torch._dynamo.config.ignore_logging_functions
  881. ):
  882. # Treat ignored functions as full no-ops
  883. self.install_guards(GuardBuilder.ID_MATCH)
  884. return IgnoredFunctionVariable(value, source=self.source)
  885. elif isinstance(value, logging.Logger):
  886. self.install_guards(GuardBuilder.TYPE_MATCH)
  887. return LoggingLoggerVariable(value, source=self.source)
  888. elif is_utils_checkpoint(value):
  889. return build_checkpoint_variable(source=self.source)
  890. elif is_invoke_subgraph(value):
  891. return build_invoke_subgraph_variable(source=self.source)
  892. elif LocalMapWrappedHigherOrderVariable.should_wrap_in_hop(value):
  893. return LocalMapWrappedHigherOrderVariable.build(source=self.source)
  894. elif isinstance(value, functools.partial):
  895. func_src = AttrSource(self.get_source(), "func")
  896. func_obj = VariableBuilder(self.tx, func_src)(value.func)
  897. args = []
  898. args_source = AttrSource(self.get_source(), "args")
  899. for i, arg in enumerate(value.args):
  900. args.append(
  901. VariableBuilder(self.tx, GetItemSource(args_source, i))(arg)
  902. )
  903. keywords = {}
  904. keywords_source = AttrSource(self.get_source(), "keywords")
  905. for k, v in value.keywords.items():
  906. if not ConstantVariable.is_literal(k):
  907. unimplemented(
  908. gb_type="functools.partial() with non-literal keyword",
  909. context=f"non-literal keyword: {k}",
  910. explanation="functools.partial() expects literal/string keywords",
  911. hints=[*graph_break_hints.USER_ERROR],
  912. )
  913. keywords[k] = VariableBuilder(
  914. self.tx, DictGetItemSource(keywords_source, k)
  915. )(v)
  916. install_guard(
  917. self.get_source().make_guard(GuardBuilder.TYPE_MATCH),
  918. keywords_source.make_guard(GuardBuilder.DICT_KEYS_MATCH),
  919. args_source.make_guard(GuardBuilder.SEQUENCE_LENGTH),
  920. )
  921. # Preserve cache_hash for SAC context_fn caching
  922. original_cache_hash = getattr(value, "cache_hash", None)
  923. return FunctoolsPartialVariable(
  924. func_obj, args, keywords, original_cache_hash=original_cache_hash
  925. )
  926. elif is_typing(value):
  927. # typing.List, typing.Mapping, etc.
  928. self.install_guards(GuardBuilder.ID_MATCH)
  929. return TypingVariable(
  930. value,
  931. source=self.source,
  932. )
  933. elif np is not None and isinstance(value, np.generic):
  934. # numpy array scalars: convert to 0D arrays
  935. return self.wrap_numpy_ndarray(np.asarray(value))
  936. elif trace_rules.is_numpy(value):
  937. assert np
  938. if istype(value, types.MethodType):
  939. # Dont guard on cython functions as they dont change ids
  940. if inspect.isfunction(value.__func__):
  941. install_guard(
  942. AttrSource(self.source, "__func__").make_guard(
  943. GuardBuilder.CLOSURE_MATCH
  944. )
  945. )
  946. elif inspect.isclass(value):
  947. self.install_guards(GuardBuilder.CLASS_MATCH)
  948. elif inspect.isfunction(value):
  949. self.install_guards(GuardBuilder.CLOSURE_MATCH)
  950. elif callable(value):
  951. self.install_guards(GuardBuilder.ID_MATCH)
  952. else:
  953. self.install_guards(GuardBuilder.TYPE_MATCH)
  954. return NumpyVariable(value, source=self.source)
  955. elif trace_rules.is_numpy_dtype(value):
  956. self.install_guards(GuardBuilder.ID_MATCH)
  957. return NumpyDTypeVariable(value, source=self.source)
  958. elif trace_rules.is_numpy_type_info(value):
  959. if isinstance(value, np.iinfo):
  960. self.install_guards(GuardBuilder.TYPE_MATCH)
  961. dt_source = AttrSource(self.source, "dtype")
  962. install_guard(dt_source.make_guard(GuardBuilder.ID_MATCH))
  963. else:
  964. self.install_guards(GuardBuilder.ID_MATCH)
  965. return ConstantLikeVariable(value, source=self.source)
  966. # NB: These can't be put in type_dispatch, they have to run later
  967. elif CollectiveFunctionRewriteVariable.can_rewrite(value):
  968. self.install_guards(GuardBuilder.CLOSURE_MATCH)
  969. return CollectiveFunctionRewriteVariable.create(
  970. self.tx,
  971. value,
  972. source=self.source,
  973. )
  974. elif istype(value, torch.autograd.function.FunctionMeta):
  975. self.install_guards(GuardBuilder.CLASS_MATCH)
  976. return AutogradFunctionVariable(
  977. value,
  978. source=self.source,
  979. )
  980. elif isinstance(value, torch.autograd.function.FunctionCtx):
  981. actual_saved_tensors = None
  982. try:
  983. # type: ignore[attr-defined]
  984. actual_saved_tensors = value.saved_tensors
  985. except RuntimeError:
  986. pass
  987. saved_tensors = []
  988. guards = [self.source.make_guard(GuardBuilder.TYPE_MATCH)]
  989. if isinstance(actual_saved_tensors, tuple):
  990. saved_tensors_source = AttrSource(self.source, "saved_tensors")
  991. guards.append(
  992. saved_tensors_source.make_guard(GuardBuilder.SEQUENCE_LENGTH)
  993. )
  994. for i, v in enumerate(actual_saved_tensors):
  995. saved_tensors.append(
  996. VariableBuilder(
  997. self.tx, GetItemSource(saved_tensors_source, i)
  998. )(v)
  999. )
  1000. install_guard(*guards)
  1001. return self.tx.output.side_effects.track_object_existing(
  1002. value,
  1003. AutogradFunctionContextVariable(
  1004. value,
  1005. source=self.source,
  1006. saved_tensors=SavedTensorBox(saved_tensors),
  1007. ),
  1008. )
  1009. elif (
  1010. isinstance(value, types.MethodType)
  1011. and istype(
  1012. getattr(value, "__self__", None), torch.autograd.function.FunctionMeta
  1013. )
  1014. and getattr(value, "__name__", "") == "apply"
  1015. and value == getattr(value.__self__, "apply", None)
  1016. ):
  1017. # handle aliased autograd function `apply` calls
  1018. install_guard(
  1019. AttrSource(self.get_source(), "__func__").make_guard(
  1020. GuardBuilder.CLOSURE_MATCH
  1021. )
  1022. )
  1023. return GetAttrVariable(
  1024. AutogradFunctionVariable(
  1025. value.__self__,
  1026. source=AttrSource(self.source, member="__self__"),
  1027. ),
  1028. "apply",
  1029. )
  1030. elif isinstance(value, torch._C._ImperativeEngine):
  1031. self.install_guards(GuardBuilder.ID_MATCH)
  1032. return AutogradEngineVariable(value, source=self.source)
  1033. elif (
  1034. value
  1035. is torch._dynamo.external_utils.FakeCompiledAutogradEngine._exec_final_callbacks_stub
  1036. ):
  1037. self.install_guards(GuardBuilder.CLOSURE_MATCH)
  1038. return LambdaVariable(
  1039. lambda: UserFunctionVariable(
  1040. torch._dynamo.external_utils.FakeCompiledAutogradEngine.exec_final_callbacks,
  1041. ).call_function(
  1042. self.tx,
  1043. (self.tx.output.side_effects.get_ca_final_callbacks_var(),),
  1044. {},
  1045. )
  1046. )
  1047. elif isinstance(value, DynamoConfigPatchProxy):
  1048. return DynamoConfigPatchVariable(value.changes)
  1049. elif isinstance(value, ErrorOnGraphBreakDecoratorContextManager):
  1050. return ErrorOnGraphBreakVariable(value.error_on_graph_break)
  1051. elif callable(value) and trace_rules.lookup_callable(value) is not None:
  1052. if trace_rules.is_callable_allowed(value):
  1053. self.tx.output.has_user_defined_allowed_in_graph = True
  1054. # type: ignore[attr-defined]
  1055. return trace_rules.lookup_callable(value).create_with_source(
  1056. value, source=self.source
  1057. )
  1058. elif np and isinstance(value, np.number):
  1059. return self.wrap_unspecialized_primitive(value)
  1060. elif isinstance(value, HigherOrderOperator):
  1061. if value is torch._higher_order_ops.invoke_subgraph:
  1062. unimplemented(
  1063. gb_type="Attempted to wrap torch._higher_order_ops.invoke_subgraph",
  1064. context="",
  1065. explanation="Directly using invoke_subgraph is not supported. Use nested_compile_region",
  1066. hints=[],
  1067. )
  1068. self.install_guards(GuardBuilder.TYPE_MATCH)
  1069. return TorchHigherOrderOperatorVariable.make(value, source=self.source)
  1070. elif isinstance(value, torch.cuda.StreamContext):
  1071. self.install_guards(GuardBuilder.ID_MATCH)
  1072. stream_source = AttrSource(self.source, "stream")
  1073. stream_var = VariableBuilder(self.tx, stream_source)(value.stream)
  1074. # type: ignore[arg-type]
  1075. return StreamContextVariable.create(self.tx, stream_var)
  1076. elif isinstance(value, torch.Stream):
  1077. # This refers to the device-agnostic torch.Stream
  1078. self.install_guards(GuardBuilder.TYPE_MATCH)
  1079. index = register_user_object(value, self.source)
  1080. stream_proxy = self.tx.output.create_proxy(
  1081. "call_function", get_external_object_by_index, (index,), {}
  1082. )
  1083. set_example_value(stream_proxy.node, value)
  1084. var = StreamVariable(
  1085. stream_proxy, value, source=self.source, user_object_index=index
  1086. )
  1087. return self.tx.output.side_effects.track_object_existing(value, var)
  1088. elif isinstance(value, (torch._C._SDPAParams)):
  1089. self.install_guards(GuardBuilder.TYPE_MATCH)
  1090. return SDPAParamsVariable.create(self.tx, value, self.source)
  1091. elif isinstance(value, torch._functorch.pyfunctorch.FuncTorchInterpreter):
  1092. self.install_guards(GuardBuilder.ID_MATCH)
  1093. return FuncTorchInterpreterVariable(value)
  1094. elif isinstance(value, torch.Event):
  1095. self.install_guards(GuardBuilder.TYPE_MATCH)
  1096. index = register_user_object(value, self.source)
  1097. event_proxy = self.tx.output.create_proxy(
  1098. "call_function",
  1099. get_external_object_by_index,
  1100. (index,),
  1101. {},
  1102. )
  1103. set_example_value(event_proxy.node, value)
  1104. return EventVariable(
  1105. event_proxy,
  1106. value,
  1107. index,
  1108. source=self.source,
  1109. )
  1110. elif (
  1111. istype(value, contextlib.nullcontext)
  1112. and inspect.getattr_static(value, "enter_result", None) is None
  1113. ):
  1114. self.install_guards(GuardBuilder.TYPE_MATCH)
  1115. return NullContextVariable(source=self.source)
  1116. elif KeyedJaggedTensorVariable.is_matching_object(value):
  1117. self.install_guards(GuardBuilder.TYPE_MATCH)
  1118. result = KeyedJaggedTensorVariable(value, source=self.source)
  1119. # TODO: this doing it manually is bad
  1120. return self.tx.output.side_effects.track_object_existing(value, result)
  1121. elif isinstance(value, torch.optim.Optimizer):
  1122. self.install_guards(GuardBuilder.ID_MATCH)
  1123. self.source = OptimizerSource(self.source)
  1124. return OptimizerVariable(value, source=self.source)
  1125. elif isinstance(value, torch.DispatchKeySet):
  1126. self.install_guards(GuardBuilder.DISPATCH_KEY_SET_MATCH)
  1127. return DispatchKeySetVariable(value)
  1128. elif WorldMetaClassVariable.is_group_member_type(value):
  1129. return WorldMetaClassVariable(value, source=self.source)
  1130. elif ProcessGroupVariable.is_process_group(value):
  1131. self.install_guards(GuardBuilder.ID_MATCH)
  1132. return ProcessGroupVariable(value, source=self.source)
  1133. elif DeviceMeshVariable.is_device_mesh(value):
  1134. # TODO: see if we need to add custom guard instead of a simple ID_MATCH
  1135. self.install_guards(GuardBuilder.EQUALS_MATCH)
  1136. return DeviceMeshVariable(value, source=self.source)
  1137. elif PlacementClassVariable.is_placement_type(value):
  1138. # TODO: see if we need to add custom guard instead of a simple ID_MATCH
  1139. self.install_guards(GuardBuilder.ID_MATCH)
  1140. return PlacementClassVariable(value, source=self.source)
  1141. elif PlacementVariable.is_placement(value):
  1142. # TODO: see if we need to add custom guard instead of a simple ID_MATCH
  1143. self.install_guards(GuardBuilder.EQUALS_MATCH)
  1144. return PlacementVariable(
  1145. value,
  1146. source=self.source,
  1147. )
  1148. elif value is OrderedSet:
  1149. self.install_guards(GuardBuilder.ID_MATCH)
  1150. return OrderedSetClassVariable()
  1151. elif (
  1152. id(value) in ITERTOOLS_TYPE_IDS
  1153. and id(value) not in ITERTOOLS_POLYFILLED_TYPE_IDS
  1154. ):
  1155. self.install_guards(GuardBuilder.CLASS_MATCH)
  1156. return ItertoolsVariable(value, source=self.source)
  1157. elif isinstance(value, _DynamicScalar):
  1158. is_int = isinstance(value, DynamicInt)
  1159. source = DynamicScalarSource(self.source, is_int)
  1160. if id(value) in self.tx.output.root_tracer.dynamic_scalar_nodes:
  1161. # If we've already seen this dynamic scalar, reuse the existing
  1162. # SymInt/SymFloat node.
  1163. node = self.tx.output.root_tracer.dynamic_scalar_nodes[id(value)]
  1164. else:
  1165. sym = self.tx.output.shape_env.create_unspecified_symbol(
  1166. value.real, # type: ignore[attr-defined]
  1167. source=source,
  1168. dynamic_dim=DimDynamic.DYNAMIC,
  1169. )
  1170. node = self.tx.output.shape_env.create_symintnode(
  1171. sym,
  1172. hint=value.real, # type: ignore[attr-defined]
  1173. source=source,
  1174. )
  1175. # Bind to graph input
  1176. sym_node_proxy = self.tx.output.root_tracer.create_graph_input(
  1177. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  1178. type(node),
  1179. node,
  1180. source=source,
  1181. )
  1182. sym_node_proxy.node.meta["grapharg"] = GraphArg(
  1183. source,
  1184. node,
  1185. False,
  1186. None,
  1187. is_tensor=False,
  1188. example_strong_ref=node,
  1189. )
  1190. sym_expr = node.node.expr
  1191. assert isinstance(sym_expr, sympy.Symbol), (
  1192. f"{sym_expr} is not a basic Symbol."
  1193. )
  1194. self.tx.output.tracked_fakes.append(TrackedFake(node, source, None))
  1195. return SymNodeVariable.create(self.tx, sym_node_proxy, node)
  1196. elif is_torch_sym(value):
  1197. # Note: this doesn't handle nested symints.
  1198. # For SymBool input, we reuse the infra for SymInt by simulating SymBool with a SymInt in dynamo.
  1199. # Concretely,
  1200. # 1. We create a SymInt in dynamo's shape_env, whose source is constructed as ConvertIntSource(self.source).
  1201. # so that guards on the SymInts can be effectively applied on the original SymBool in user program.
  1202. # 2. We create a SymBool based on the SymInt in dynamo's ShapeEnv. Because the original user program
  1203. # depends on the value being a SymBool. This allows dynamo to interpret the user's program correctly.
  1204. source = (
  1205. self.source
  1206. if isinstance(value, torch.SymInt)
  1207. else ConvertIntSource(self.source)
  1208. )
  1209. new_symint = None
  1210. if value.node.has_hint():
  1211. new_symint = (
  1212. self.tx.output.shape_env.create_unspecified_symint_and_symbol(
  1213. int(value.node.hint),
  1214. source,
  1215. dynamic_dim=DimDynamic.DYNAMIC,
  1216. )
  1217. )
  1218. else:
  1219. if isinstance(value, torch.SymBool):
  1220. # We need to create an unbacked symint to replace the unbacked symbool.
  1221. new_symint = self.tx.output.shape_env.create_unbacked_symint()
  1222. else:
  1223. # TODO (yidi): we need to figure out a way to propagate the guards
  1224. # we accumulated when tracing the subggraph to outer shape_env. For normal symints,
  1225. # this is automatically done by evaluating the guards once but this
  1226. # will cause data-dependent error when we evaluate the outer unbacked symints.
  1227. # The test case that triggers this graph break is test_cond_unbacked_symint_closure
  1228. unimplemented(
  1229. gb_type="Attempted to wrap unbacked SymInt",
  1230. context="",
  1231. explanation="Unbacked SymInt input is not supported yet.",
  1232. hints=[*graph_break_hints.SUPPORTABLE],
  1233. )
  1234. assert new_symint is not None
  1235. sym_node_proxy = self.tx.output.root_tracer.create_graph_input(
  1236. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  1237. type(new_symint),
  1238. new_symint,
  1239. source=source,
  1240. )
  1241. sym_node_proxy.node.meta["grapharg"] = GraphArg(
  1242. source,
  1243. new_symint,
  1244. False,
  1245. None,
  1246. is_tensor=False,
  1247. example_strong_ref=new_symint,
  1248. )
  1249. # We bind the new_symint to graph input.
  1250. sym_expr = new_symint.node.expr
  1251. assert isinstance(sym_expr, sympy.Symbol), (
  1252. f"{sym_expr} is not a basic Symbol."
  1253. )
  1254. self.tx.output.tracked_fakes.append(TrackedFake(new_symint, source, None))
  1255. tracing_symint = (
  1256. new_symint if isinstance(value, torch.SymInt) else new_symint == 1
  1257. ) # cast it back to symbool for tracing
  1258. return SymNodeVariable(sym_node_proxy, tracing_symint)
  1259. elif isinstance(value, (JITFunction, Autotuner)):
  1260. self.install_guards(GuardBuilder.ID_MATCH)
  1261. return TritonKernelVariable(
  1262. value,
  1263. None, # No kernel idx provided
  1264. None, # No grid provided
  1265. source=self.source,
  1266. )
  1267. elif value is create_1d_tma_descriptor:
  1268. return CreateTMADescriptorExperimentalVariable(rank=1)
  1269. elif value is create_2d_tma_descriptor:
  1270. return CreateTMADescriptorExperimentalVariable(rank=2)
  1271. elif value is TensorDescriptor.from_tensor:
  1272. return CreateTMADescriptorStableVariable()
  1273. elif value is set_allocator:
  1274. return TritonSetAllocatorSkipVariable(value)
  1275. elif isinstance(value, torch.amp.autocast_mode.autocast):
  1276. self.install_guards(GuardBuilder.ID_MATCH)
  1277. return AutocastModeVariable(
  1278. target_values=[
  1279. value.device,
  1280. value.fast_dtype,
  1281. value._enabled,
  1282. value._cache_enabled,
  1283. ],
  1284. source=self.source,
  1285. )
  1286. elif TorchCtxManagerClassVariable.is_matching_cls(value):
  1287. if inspect.isclass(value):
  1288. self.install_guards(GuardBuilder.CLASS_MATCH)
  1289. elif inspect.isfunction(value):
  1290. self.install_guards(GuardBuilder.CLOSURE_MATCH)
  1291. return TorchCtxManagerClassVariable(value, source=self.source)
  1292. elif inspect.getattr_static(value, "__script_if_tracing_wrapper", False):
  1293. self.install_guards(GuardBuilder.TYPE_MATCH)
  1294. return WrapperUserFunctionVariable(
  1295. value, "__original_fn", source=self.source
  1296. )
  1297. elif is_lru_cache_wrapped_function(value):
  1298. self.install_guards(GuardBuilder.TYPE_MATCH)
  1299. return WrapperUserFunctionVariable(value, "__wrapped__", source=self.source)
  1300. elif value is sys.exc_info or (
  1301. sys.version_info >= (3, 11) and value is sys.exception
  1302. ):
  1303. return SysFunctionVariable(value, source=self.source)
  1304. elif is_function_or_wrapper(value) and inspect.getattr_static(
  1305. value, "_torchdynamo_inline", False
  1306. ):
  1307. self.install_guards(GuardBuilder.TYPE_MATCH)
  1308. return WrapperUserFunctionVariable(
  1309. value, "_torchdynamo_inline", source=self.source
  1310. )
  1311. elif value is functools.wraps:
  1312. self.install_guards(GuardBuilder.ID_MATCH)
  1313. return FunctoolsWrapsVariable(value, source=self.source)
  1314. elif value is collections.namedtuple:
  1315. self.install_guards(GuardBuilder.ID_MATCH)
  1316. return CollectionsNamedTupleFunction(value, source=self.source)
  1317. elif isinstance(
  1318. value, types.BuiltinMethodType
  1319. ) and BuiltinMethodVariable.is_supported_builtin_method(value):
  1320. self.install_guards(GuardBuilder.ID_MATCH)
  1321. return BuiltinMethodVariable(value, source=self.source)
  1322. elif is_function(value) and value in (float.fromhex, float.hex):
  1323. self.install_guards(GuardBuilder.ID_MATCH)
  1324. return GetAttrVariable(
  1325. BuiltinVariable(float, source=self.source),
  1326. value.__name__,
  1327. )
  1328. elif is_function_or_wrapper(value):
  1329. value, attr_name = unwrap_with_attr_name_if_wrapper(value)
  1330. # For these wrappers, Dynamo points to the wrapped function,
  1331. # so source needs to be updated as well.
  1332. if attr_name is not None:
  1333. self.source = AttrSource(self.source, attr_name)
  1334. # type: ignore[attr-defined]
  1335. return trace_rules.lookup(value).create_with_source(
  1336. value, source=self.source
  1337. )
  1338. elif value is random.Random:
  1339. self.install_guards(GuardBuilder.ID_MATCH)
  1340. return RandomClassVariable(source=self.source)
  1341. elif istype(value, random.Random) and RandomVariable.is_supported_random_obj(
  1342. value
  1343. ):
  1344. self.install_guards(GuardBuilder.TYPE_MATCH)
  1345. result = RandomVariable(value, source=self.source)
  1346. self.tx.output.side_effects.track_mutable(value, result)
  1347. return result
  1348. # Don't use istype, since some python modules are not subclasses of types.ModuleType directly.
  1349. # E.g, type(torch.ops) -> <class 'torch._ops._Ops'>,
  1350. # type(torch.backends.cudnn) -> <class 'torch.backends.cudnn.CudnnModule'>
  1351. elif isinstance(value, (types.ModuleType, replay_record.DummyModule)):
  1352. self.install_guards(GuardBuilder.MODULE_MATCH)
  1353. result = PythonModuleVariable(
  1354. # type: ignore[arg-type]
  1355. value,
  1356. source=self.source,
  1357. )
  1358. self.tx.output.side_effects.track_object_existing(value, result)
  1359. return result
  1360. elif isinstance(value, types.MethodType) and isinstance(
  1361. value.__self__, (torch.nn.Module, torch.utils._pytree.TreeSpec)
  1362. ):
  1363. # don't let MethodTypes fall through to UserDefinedObject,
  1364. # which doesn't support 'CALL_FUNCTION'
  1365. # TODO(whc): Why do we limit this to methods on NNModules?
  1366. # I don't have a good reason for this, but it preserves the existing behavior
  1367. # for MBartForConditionalGeneration, which generates many graph breaks and OOMs otherwise.
  1368. # I suspect we probably want to relax this check and dig deeper there.
  1369. # In order to construct a MethodVariable in Dynamo, we start with an actual method obj from python,
  1370. # but need to separately wrap its underlying `__func__` and its `self` argument. We wrap `self` here
  1371. # and then `__func__` gets wrapped inside UserMethodVariable.
  1372. self_obj = VariableBuilder(
  1373. self.tx, source=AttrSource(self.source, "__self__")
  1374. )(value.__self__)
  1375. assert self_obj and isinstance(self_obj, VariableTracker), (
  1376. "Failed to produce a valid self obj"
  1377. )
  1378. return UserMethodVariable(
  1379. value.__func__,
  1380. self_obj,
  1381. source=self.source,
  1382. )
  1383. elif isinstance(value, types.GetSetDescriptorType):
  1384. # GetSet descriptors are C functions attached to an attribute lookup
  1385. # using PyGetSetDef. Python, on attribute lookup, can decide to
  1386. # create a new object on the fly, and therefore the `id` of the
  1387. # descriptors is not guaranteed to be same for different attribute
  1388. # accesses. Since these are unlikely to change during the program
  1389. # execution, we can skip guarding on them.
  1390. return GetSetDescriptorVariable(value)
  1391. elif isinstance(value, types.MethodWrapperType):
  1392. # Method-wrappers are written in C, and they are not guaranteed to
  1393. # return the same object on attribute lookup. Therefore, we cannot
  1394. # insert a ID_MATCH guard here. method-wrappers are very
  1395. # unlikely to change, so its ok to skip the guard here.
  1396. return MethodWrapperVariable(value)
  1397. elif issubclass(type(value), type) and issubclass(value, BaseException):
  1398. # match user defined exceptions
  1399. self.install_guards(GuardBuilder.ID_MATCH)
  1400. return UserDefinedExceptionClassVariable(value)
  1401. elif issubclass(type(value), type):
  1402. if value in (
  1403. torch.utils.hooks.BackwardHook,
  1404. torch.nn.Parameter,
  1405. torch.nn.Buffer,
  1406. ):
  1407. # TODO(jansel): combine this case with the one above
  1408. # type: ignore[attr-defined]
  1409. return trace_rules.lookup(value).create_with_source(
  1410. value, source=self.source
  1411. )
  1412. if value is torch.autograd._unsafe_preserve_version_counter:
  1413. self.install_guards(GuardBuilder.CLASS_MATCH)
  1414. return PreserveVersionContextVariable.constructor(self.tx)
  1415. if (
  1416. # `value` must be a strict subclass of `torch.Tensor`
  1417. issubclass(value, torch.Tensor)
  1418. and value is not torch.Tensor
  1419. # `TensorSubclassVariable` is not for subclass that overrides
  1420. # `torch_dispatch`.
  1421. and value.__torch_dispatch__ is torch.Tensor.__torch_dispatch__
  1422. # `TensorSubclassVariable` would lead to construction of
  1423. # `TensorWithTFOverrideVariable`, but we don't want that for
  1424. # traceable wrapper subclasses (we wrap those subclass instances
  1425. # into `TensorVariable`).
  1426. and not is_traceable_wrapper_subclass_type(value)
  1427. ):
  1428. return TensorSubclassVariable(value, source=self.source)
  1429. if not is_from_closure_source(self.source):
  1430. # For closure source, the variable comes from LOAD_SUPER_ATTR,
  1431. # which calls self.__class__. This is internal Cpython
  1432. # implementation, and it is rare for the user to modify
  1433. # self.__class__ manually.
  1434. # For other cases, this is a userdefined class, so install an
  1435. # ID_MATCH even if its a global variable.
  1436. self.install_guards(GuardBuilder.CLASS_MATCH)
  1437. if is_opaque_type(value):
  1438. return OpaqueObjectClassVariable(
  1439. value,
  1440. source=self.source,
  1441. )
  1442. if isinstance(value, type) and issubclass(value, enum.Enum):
  1443. return UserDefinedEnumClassVariable(
  1444. value,
  1445. source=self.source,
  1446. )
  1447. return UserDefinedClassVariable(
  1448. value,
  1449. source=self.source,
  1450. )
  1451. elif TorchScriptObjectVariable.is_matching_cls(type(value)):
  1452. from ..source import (
  1453. FlattenScriptObjectSource,
  1454. ScriptObjectQualifiedNameSource,
  1455. )
  1456. # type: ignore[arg-type]
  1457. if torch._library.fake_class_registry.tracing_with_real(value):
  1458. proxy = self.tx.output.root_tracer.create_graph_input(
  1459. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  1460. type(value),
  1461. value,
  1462. source=self.source,
  1463. )
  1464. # setting is_unspecialized=False to not insert a as_tensor call in reconstruct by default
  1465. # setting example to be real value because these example values will be used
  1466. # as example_inputs for user compiler.
  1467. proxy.node.meta["grapharg"] = GraphArg(
  1468. self.source,
  1469. value,
  1470. False,
  1471. None,
  1472. False,
  1473. value, # type: ignore[arg-type]
  1474. )
  1475. return TorchScriptObjectVariable.create(
  1476. proxy,
  1477. value,
  1478. source=self.source,
  1479. )
  1480. if is_opaque_value_type(type(value)):
  1481. # Value-type: guard on equality (will use __eq__)
  1482. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  1483. elif is_opaque_reference_type(type(value)):
  1484. # Reference-type: guard only on type, and registered guard_fn
  1485. self.install_guards(GuardBuilder.TYPE_MATCH)
  1486. self.install_guards(GuardBuilder.OPAQUE_OBJ_GUARD_FN_MATCH)
  1487. elif not hasattr(value, "__obj_flatten__"):
  1488. # This exists to allow a smoother transition.
  1489. # The implications are:
  1490. # The script objects won't be tracked as proxies.
  1491. # Methods on these objects won't show up in the graph.
  1492. # The original script object might be mutated.
  1493. return self.wrap_user_defined(value)
  1494. else:
  1495. # Install the guards on the fully qualified name of the script object
  1496. LazyVariableTracker.realize_all(
  1497. VariableBuilder(
  1498. self.tx, ScriptObjectQualifiedNameSource(self.source)
  1499. )(
  1500. value._type().qualified_name() # type: ignore[attr-defined]
  1501. )
  1502. )
  1503. # Install the guards on the content of the script object by setting the source
  1504. # to be FlattenScriptObjectSource, which calls __obj_flatten__() to get the contents.
  1505. LazyVariableTracker.realize_all(
  1506. VariableBuilder(self.tx, FlattenScriptObjectSource(self.source))(
  1507. value.__obj_flatten__()
  1508. )
  1509. )
  1510. fake_script_obj = torch._library.fake_class_registry.maybe_to_fake_obj(
  1511. self.tx.output.fake_mode, value
  1512. )
  1513. if is_opaque_value_type(type(value)) and not should_hoist(type(value)):
  1514. proxy = value
  1515. else:
  1516. proxy = self.tx.output.root_tracer.create_graph_input(
  1517. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  1518. type(value),
  1519. fake_script_obj,
  1520. source=self.source,
  1521. )
  1522. # setting is_unspecialized=False to not insert a as_tensor call in reconstruct by default
  1523. # setting example to be real value because these example values will be used
  1524. # as example_inputs for user compiler.
  1525. proxy.node.meta["grapharg"] = GraphArg(
  1526. self.source,
  1527. value, # type: ignore[arg-type]
  1528. False,
  1529. None,
  1530. False,
  1531. fake_script_obj, # type: ignore[arg-type]
  1532. )
  1533. return TorchScriptObjectVariable.create(
  1534. proxy, # pyrefly: ignore[bad-argument-type]
  1535. fake_script_obj,
  1536. source=self.source,
  1537. )
  1538. elif (
  1539. isinstance(value, (dict, collections.OrderedDict))
  1540. and type(value).__new__ is dict.__new__
  1541. ):
  1542. # Construct a dict_vt that will reside inside the UserDefinedDictVariable
  1543. self.install_guards(GuardBuilder.TYPE_MATCH)
  1544. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  1545. # Guard on the key order
  1546. self.tx.output.guard_on_key_order.add(self.source)
  1547. # We need all the keys to be hashable. We do this within the
  1548. # _HashableTracker class in dicts.py
  1549. def build_key_value(
  1550. i: Any, k: Any, v: Any
  1551. ) -> tuple[VariableTracker, VariableTracker]:
  1552. base = self.get_source()
  1553. source_key = ConstDictKeySource(base, i)
  1554. key = LazyVariableTracker.create(k, source_key)
  1555. source_value = DictSubclassGetItemSource(base, source_key)
  1556. res_value = LazyVariableTracker.create(v, source_value)
  1557. return key, res_value
  1558. # Ensure that we call dict.keys and not value.keys (which can call
  1559. # overridden keys method). In the C++ guards, we relied on
  1560. # PyDict_Next to traverse the dictionary, which uses the internal
  1561. # data structure and does not call the overridden keys method.
  1562. result = dict(
  1563. build_key_value(i, k, v)
  1564. for i, (k, v) in enumerate(get_items_from_dict(value))
  1565. )
  1566. dict_vt = ConstDictVariable(
  1567. # pyrefly: ignore[bad-argument-type]
  1568. result,
  1569. user_cls=(
  1570. collections.OrderedDict
  1571. if isinstance(value, collections.OrderedDict)
  1572. else dict
  1573. ),
  1574. mutation_type=ValueMutationExisting(),
  1575. source=self.source,
  1576. )
  1577. # Force this to reconstruct on mutation to keep the reconstruction
  1578. # bytecode simple
  1579. dict_vt.should_reconstruct_all = True
  1580. result = UserDefinedDictVariable(value, dict_vt=dict_vt, source=self.source)
  1581. return self.tx.output.side_effects.track_object_existing(value, result)
  1582. elif isinstance(value, tuple):
  1583. self.install_guards(GuardBuilder.TYPE_MATCH)
  1584. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  1585. # NB - Be careful in not triggering user code. Guards also work on
  1586. # the underlying tuple data structure.
  1587. output = [
  1588. LazyVariableTracker.create(
  1589. tuple.__getitem__(value, i),
  1590. source=GetItemSource(self.get_source(), i),
  1591. )
  1592. for i in range(tuple.__len__(value))
  1593. ]
  1594. tuple_vt = TupleVariable(
  1595. output, # type: ignore[arg-type]
  1596. source=self.source,
  1597. mutation_type=ValueMutationExisting(),
  1598. )
  1599. result = UserDefinedTupleVariable(
  1600. value, tuple_vt=tuple_vt, source=self.source
  1601. )
  1602. return self.tx.output.side_effects.track_object_existing(value, result)
  1603. elif isinstance(value, list):
  1604. self.install_guards(GuardBuilder.TYPE_MATCH)
  1605. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  1606. # NB - Be careful in not triggering user code. Guards also work on
  1607. # the underlying list data structure.
  1608. output = [
  1609. LazyVariableTracker.create(
  1610. list.__getitem__(value, i),
  1611. source=ListGetItemSource(self.get_source(), i),
  1612. )
  1613. for i in range(list.__len__(value))
  1614. ]
  1615. list_vt = ListVariable(
  1616. output, # type: ignore[arg-type]
  1617. source=self.source,
  1618. mutation_type=ValueMutationExisting(),
  1619. )
  1620. result = UserDefinedListVariable(value, list_vt=list_vt, source=self.source)
  1621. return self.tx.output.side_effects.track_object_existing(value, result)
  1622. elif isinstance(value, (set, frozenset)):
  1623. self.install_guards(GuardBuilder.TYPE_MATCH)
  1624. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  1625. L = list(dict.fromkeys(value))
  1626. output = [
  1627. LazyVariableTracker.create(
  1628. list.__getitem__(L, i),
  1629. source=NonSerializableSetGetItemSource(self.get_source(), i),
  1630. )
  1631. for i in range(list.__len__(L))
  1632. ]
  1633. if isinstance(value, set):
  1634. set_vt_cls = SetVariable
  1635. else:
  1636. assert isinstance(value, frozenset)
  1637. set_vt_cls = FrozensetVariable
  1638. set_vt = set_vt_cls(
  1639. output, source=self.source, mutation_type=ValueMutationExisting()
  1640. )
  1641. result = UserDefinedSetVariable(value, set_vt=set_vt, source=self.source)
  1642. return self.tx.output.side_effects.track_object_existing(value, result)
  1643. elif issubclass(type(value), MutableMapping):
  1644. self.install_guards(GuardBuilder.TYPE_MATCH)
  1645. result = MutableMappingVariable(value, source=self.source)
  1646. return self.tx.output.side_effects.track_object_existing(value, result)
  1647. elif is_frozen_dataclass(value):
  1648. self.install_guards(GuardBuilder.TYPE_MATCH)
  1649. result = FrozenDataClassVariable.create(self.tx, value, source=self.source)
  1650. return self.tx.output.side_effects.track_object_existing(value, result)
  1651. elif isinstance(value, dict_keys):
  1652. if all(ConstantVariable.is_literal(k) for k in value):
  1653. # If the dict_keys object is passed from outside the compile region, it must either be passed along with
  1654. # the corresponding dict object or treated as a set (when only the keys are passed into the compiled region).
  1655. # - If it is passed along with the dict, the dict object itself is already guarded.
  1656. # - If only the dict_keys object is passed, we add EQUALS_MATCH and SEQUENCE_LENGTH guards
  1657. # to ensure it remains unchanged across multiple runs.
  1658. items = [SourcelessBuilder.create(self.tx, v) for v in value]
  1659. install_guard(
  1660. self.get_source().make_guard(GuardBuilder.SEQUENCE_LENGTH),
  1661. self.get_source().make_guard(GuardBuilder.EQUALS_MATCH),
  1662. )
  1663. return DictKeySetVariable(items, source=self.source)
  1664. else:
  1665. unimplemented(
  1666. gb_type="non-const keys in dict_keys",
  1667. context=f"non-const keys: {[k for k in value if not ConstantVariable.is_literal(k)]}",
  1668. explanation="Dynamo expects dict_keys keys to be constants.",
  1669. hints=[
  1670. "Ensure your dict_keys keys are constants (e.g. int, float, strings)",
  1671. ],
  1672. )
  1673. elif IntWrapperVariable.is_matching_object(value):
  1674. from torch.export.dynamic_shapes import _DimHintType
  1675. if value.dynamism is None or value.dynamism.type == _DimHintType.STATIC:
  1676. return self.wrap_symint(value.val)
  1677. elif value.dynamism.type == _DimHintType.DYNAMIC:
  1678. log.debug(
  1679. "%s marked %s via IntWrapper",
  1680. self.source.name,
  1681. DimDynamic.DYNAMIC,
  1682. )
  1683. return self.wrap_symint(
  1684. value.val,
  1685. dynamism=DimDynamic.DYNAMIC,
  1686. context=SymIntSymbolicContext(
  1687. constraint=RelaxedUnspecConstraint(warn_only=False)
  1688. ),
  1689. )
  1690. elif value.dynamism.type == _DimHintType.AUTO:
  1691. log.debug(
  1692. "%s marked %s via IntWrapper",
  1693. self.source.name,
  1694. DimDynamic.DYNAMIC,
  1695. )
  1696. return self.wrap_symint(value.val, dynamism=DimDynamic.DYNAMIC)
  1697. else:
  1698. raise RuntimeError(f"Undefined dynamism {value.dynamism}")
  1699. elif istype(value, object):
  1700. self.install_guards(GuardBuilder.TYPE_MATCH)
  1701. return ObjectVariable(value, source=self.source)
  1702. else:
  1703. return self.wrap_user_defined(value)
  1704. def wrap_user_defined(self, value: Any) -> VariableTracker:
  1705. self.install_guards(GuardBuilder.TYPE_MATCH)
  1706. if InspectVariable.is_matching_object(value):
  1707. # Skip guards on inspect related variable trackers because they are
  1708. # not important for recompiles (something else will also change to
  1709. # cause recompiles) and can cause a large number of OBJECT_ALIASING
  1710. # guards.
  1711. result = InspectVariable(value, source=SkipGuardSource(self.source))
  1712. else:
  1713. result = UserDefinedObjectVariable(value, source=self.source)
  1714. if not SideEffects.cls_supports_mutation_side_effects(type(value)):
  1715. # don't allow STORE_ATTR mutation with custom __setattr__
  1716. return result
  1717. return self.tx.output.side_effects.track_object_existing(value, result)
  1718. def wrap_listlike(
  1719. self, value: Union[tuple[Any, ...], list[Any], odict_values, NamedTuple]
  1720. ) -> VariableTracker:
  1721. if config.specialize_int and type(value) is torch.Size:
  1722. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  1723. return ConstantVariable.create(value=value)
  1724. # One can index a tensor with a list/tuple. Therefore, we need to
  1725. # have a stricter match.
  1726. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  1727. # Tuples are immutable objects, so we should mark its items static. This
  1728. # avoids wrapping of tuple items as symints. This helps for nn module
  1729. # attributes like conv2d strides, dilations.
  1730. if (
  1731. istype(value, tuple)
  1732. and all(ConstantVariable.is_literal(item) for item in value)
  1733. and self.source.guard_source.is_unspecialized_nn_module()
  1734. ):
  1735. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  1736. return TupleVariable([ConstantVariable.create(item) for item in value])
  1737. output = [
  1738. LazyVariableTracker.create(
  1739. item,
  1740. source=GetItemSource(self.get_source(), i),
  1741. )
  1742. for i, item in enumerate(value)
  1743. ]
  1744. maybe_gm = self.tx.output.local_scope.get("self")
  1745. if isinstance(
  1746. self.source, LocalSource
  1747. ) and self.source.local_name in get_locals_to_steal(maybe_gm):
  1748. # The input tensor list to dynamo from compiled autograd may contain activations
  1749. # which are freed as they are used in inductor. Dynamo's default behavior is to
  1750. # lift all tensors to the graph inputs, but this will cause dynamo to hold an
  1751. # extra reference to the activation tensors and increase peak memory usage.
  1752. # To allow freeing ASAP, we keep the list as graph argument to the dynamo output
  1753. # graph, and unpack it locally.
  1754. # e.g. instead of `def forward(self, L_inputs_0_, L_inputs_1_, ...):`, we have
  1755. # `def forward(self, L_inputs_):`
  1756. source = self.source
  1757. assert isinstance(value, list)
  1758. tensor_list_proxy = self.tx.output.root_tracer.create_graph_input(
  1759. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  1760. type(value),
  1761. value,
  1762. source=source,
  1763. )
  1764. tensor_list_proxy.node.meta["steal_arg"] = True
  1765. list_variable = wrap_fx_proxy_cls(
  1766. target_cls=TensorVariable,
  1767. tx=self.tx,
  1768. proxy=tensor_list_proxy,
  1769. example_value=value,
  1770. subclass_type=None,
  1771. source=source,
  1772. )
  1773. # Apply relevant logic from `VariableTracker.build(value[i])`
  1774. # (except for the `create_graph_input` stuff).
  1775. guards = []
  1776. # type: ignore[attr-defined]
  1777. for i, tensor_variable in enumerate(list_variable.items):
  1778. source_i = GetItemSource(base=source, index=i, index_is_slice=False)
  1779. # access unpacked tensor from this list instead of from a lifted arg
  1780. self.tx.output.input_source_to_var[source_i] = tensor_variable
  1781. tensor_variable.proxy.node.meta["tensor_dict"] = _extract_tensor_dict(
  1782. value[i]
  1783. )
  1784. guard = functools.partial(
  1785. GuardBuilder.TENSOR_MATCH, value=TensorWeakRef(value[i])
  1786. )
  1787. guards.append(source_i.make_guard(guard))
  1788. install_guard(*guards, skip=1)
  1789. grapharg = GraphArg(
  1790. source,
  1791. value,
  1792. pass_arg_as_tensor=False,
  1793. fake_tensor=None,
  1794. is_tensor=False,
  1795. )
  1796. tensor_list_proxy.node.meta["grapharg"] = grapharg
  1797. # The following is very important for maintaining the "python object
  1798. # <==> variable tracker" 1-to-1 mapping, which is mainly handled via
  1799. # `side_effects`. Note that constructing `tensor_variable` above
  1800. # already adds it to graph arg, but we never registered it with
  1801. # `side_effects`. The preemptive `realize` calls here basically
  1802. # does that registration (at the end of `self.__call__`).
  1803. #
  1804. # A slightly cleaner alternative is to register the
  1805. # `tensor_variable`s above with `side_effects` directly, and just
  1806. # return the `list_variable`, but that breaks some tensor-subclass
  1807. # related tests like `test_inputs_aliasing_bytecode_stack_restore`,
  1808. # because `tensor_variable` is constructed via
  1809. # `handle_traced_output`, which doesn't really expect/handle tensor
  1810. # subclass.
  1811. #
  1812. # Eventually, we expect to fix remove all of these by having Dynamo
  1813. # auto-boxing inputs to the compiled graph, see
  1814. # https://github.com/pytorch/pytorch/issues/153701.
  1815. for vt in output:
  1816. vt.realize()
  1817. # type: ignore[arg-type]
  1818. result = BaseListVariable.cls_for_instance(value)(output, source=self.source)
  1819. if istype(value, (list, collections.deque)):
  1820. return self.tx.output.side_effects.track_mutable(value, result)
  1821. return result
  1822. def wrap_tuple_iterator(self, value: tuple_iterator) -> VariableTracker:
  1823. self.install_guards(GuardBuilder.TUPLE_ITERATOR_LEN)
  1824. output = [
  1825. VariableBuilder(self.tx, TupleIteratorGetItemSource(self.get_source(), i))(
  1826. tuple_iterator_getitem(value, i)
  1827. )
  1828. for i in range(tuple_iterator_len(value))
  1829. ]
  1830. result = TupleIteratorVariable(output, source=self.source)
  1831. return self.tx.output.side_effects.track_mutable(value, result)
  1832. def wrap_range_iterator(self, value: range_iterator) -> VariableTracker:
  1833. self.install_guards(GuardBuilder.RANGE_ITERATOR_MATCH)
  1834. # Get all the values from the range iterator; no need to install guards
  1835. # on items since `RANGE_ITERATOR_MATCH` guarantees the same items.
  1836. items = [ConstantVariable.create(v) for v in copy.deepcopy(value)]
  1837. result = ListIteratorVariable(items, source=self.source)
  1838. return self.tx.output.side_effects.track_mutable(value, result)
  1839. def wrap_slice_range(self, value: slice | range) -> SliceVariable | RangeVariable:
  1840. items = [
  1841. VariableBuilder(self.tx, AttrSource(self.get_source(), k))(
  1842. getattr(value, k)
  1843. )
  1844. for k in ("start", "stop", "step")
  1845. ]
  1846. self.install_guards(GuardBuilder.TYPE_MATCH)
  1847. if isinstance(value, slice):
  1848. return SliceVariable(items, self.tx, source=self.source)
  1849. else:
  1850. return RangeVariable(items, source=self.source)
  1851. def mark_static_input(self, value: torch.Tensor, guard: bool) -> None:
  1852. from ..decorators import mark_static_address
  1853. static_inputs_log.debug(
  1854. "Marking static input %s, id: %s)", self.source.name, id(value)
  1855. )
  1856. mark_static_address(value, guard=guard)
  1857. # Check if we've seen this tensor before and update graph metadata if needed
  1858. # As long as this runs before AOT this is sound
  1859. if value in self.tx.output.side_effects:
  1860. var = self.tx.output.side_effects[value]
  1861. # type: ignore[attr-defined]
  1862. var.proxy.node.meta["tensor_dict"]["_dynamo_static_input_type"] = (
  1863. # type: ignore[attr-defined]
  1864. value._dynamo_static_input_type
  1865. )
  1866. def wrap_module(self, value: torch.nn.Module) -> VariableTracker:
  1867. from ..eval_frame import OptimizedModule
  1868. if len(value.__dict__) == 0:
  1869. unimplemented(
  1870. gb_type="Uninitialized nn.Module",
  1871. context=typestr(value),
  1872. explanation=f"Attempted to trace an uninitialized nn.Module of type {typestr(value)}.",
  1873. hints=[
  1874. *graph_break_hints.USER_ERROR,
  1875. "Ensure your nn.Module instance has called `super().__init__()`.",
  1876. ],
  1877. )
  1878. if istype(value, OptimizedModule):
  1879. # Check if the optimized module was disabled
  1880. if inspect.getattr_static(value.forward, "_torchdynamo_disable", False):
  1881. # This bytecode is mostly of kind LOAD_ATTR or LOAD_METHOD. If
  1882. # we graph break here, Dynamo does not know how to create
  1883. # continuation functions for such bytecodes. So, we delay the
  1884. # graph break to CALL_FUNCTION.
  1885. msg = inspect.getattr_static(
  1886. value.forward, "_torchdynamo_disable_msg", None
  1887. )
  1888. return DelayGraphBreakVariable(
  1889. source=self.source,
  1890. msg=f"Optimized `nn.Module` is wrapped with `torch.compiler.disable` (reason: {msg})",
  1891. )
  1892. self.install_guards(GuardBuilder.TYPE_MATCH)
  1893. self.source = AttrSource(self.source, "_orig_mod")
  1894. return self.wrap_module(value._orig_mod)
  1895. if (
  1896. isinstance(value, (torch.nn.RNN, torch.nn.GRU, torch.nn.LSTM))
  1897. and not config.allow_rnn
  1898. ):
  1899. unimplemented(
  1900. gb_type="Attempted to wrap RNN, GRU, or LSTM",
  1901. context=str(value),
  1902. explanation="Dynamo does not support RNN, GRU, or LSTM.",
  1903. hints=[
  1904. "Set torch._dynamo.config.allow_rnn=True to enable experimental support for RNN, GRU, and LSTM in Dynamo",
  1905. *graph_break_hints.SUPPORTABLE,
  1906. ],
  1907. )
  1908. if getattr(value, "_is_fsdp_managed_module", False):
  1909. # See note [Dynamo treats FSDP wrapped modules as UnspecializedNNModule]
  1910. # in fully_sharded_data_parallel.py for more information
  1911. # we can't do this assert inside FSDP constructor,
  1912. # since we don't know yet whether dynamo will be used
  1913. if not getattr(value, "_fsdp_use_orig_params", False):
  1914. unimplemented(
  1915. gb_type="FSDP with use_orig_params=False",
  1916. context="",
  1917. explanation="Dynamo only supports FSDP with use_orig_params=True",
  1918. hints=[],
  1919. )
  1920. # Note on FSDP guarding
  1921. # Eager FSDP already assumes (requires, but without enforcement)
  1922. # that users don't mutate their model parameters/structure after
  1923. # FSDP wrapping, because FSDP wouldn't notice or update its
  1924. # FlatParams.
  1925. #
  1926. # Therefore, torch.compile can skip guarding on params or submodule
  1927. # structure of fsdp_managed modules, by using FSDPNNModuleSource as
  1928. # the guard source. This behavior is gated on
  1929. # config.skip_fsdp_guards.
  1930. self.install_guards(GuardBuilder.TYPE_MATCH)
  1931. result = FSDPManagedNNModuleVariable(value, source=self.get_source())
  1932. if not SideEffects.cls_supports_mutation_side_effects(type(value)):
  1933. # don't allow STORE_ATTR mutation with custom __setattr__
  1934. return result
  1935. return self.tx.output.side_effects.track_object_existing(value, result)
  1936. elif mutation_guard.is_dynamic_nn_module(value, self.tx.export):
  1937. # created dynamically, don't specialize on it
  1938. # Note [Tracing a torch.compiled function]
  1939. # when make_fx tracing a compiled function, we need
  1940. if isinstance(value, torch.fx.experimental.proxy_tensor._AttrProxy):
  1941. # type: ignore[attr-defined]
  1942. value = value.get_base()
  1943. self.source = AttrProxySource(self.source)
  1944. if torch._dynamo.config.inline_inbuilt_nn_modules:
  1945. freezing = is_parameter_freezing()
  1946. # Guard against the case where user may overwrite named parameters
  1947. # / named buffers
  1948. # NOTE: This is not likely to happen but worth guarding to avoid
  1949. # exception
  1950. if (
  1951. callable(value.named_parameters)
  1952. # type: ignore[attr-defined]
  1953. and value.named_parameters.__func__
  1954. is og_module_named_parameters_fn_ptr
  1955. ):
  1956. try: # catch TypeErrors in named_parameters() from unserializable nn modules
  1957. # type: ignore[attr-defined]
  1958. for _, p in value.named_parameters():
  1959. self.mark_static_input(p, guard=freezing)
  1960. except TypeError as e:
  1961. raise_observed_exception(type(e), self.tx, args=list(e.args))
  1962. if (
  1963. callable(value.named_buffers)
  1964. # type: ignore[attr-defined]
  1965. and value.named_buffers.__func__ is og_module_named_buffers_fn_ptr
  1966. ):
  1967. try: # catch TypeErrors in named_parameters() from unserializable nn modules
  1968. # type: ignore[attr-defined]
  1969. for _, b in value.named_buffers():
  1970. self.mark_static_input(b, guard=freezing)
  1971. except TypeError as e:
  1972. raise_observed_exception(type(e), self.tx, args=list(e.args))
  1973. if freezing:
  1974. # we need to add the module to tracing context
  1975. # in order to allow its params to get invalidated
  1976. # this will get cleaned up once compile ends
  1977. self.tx.output.nn_modules[self.name] = value
  1978. if (
  1979. value.__module__.startswith(("torch.nn.modules", "torch.ao."))
  1980. and not value.__module__.startswith("torch.nn.modules.container")
  1981. ) or getattr(value.__class__, "_dynamo_marked_static", False):
  1982. new_source = self.source
  1983. if config.inline_inbuilt_nn_modules and (
  1984. not self.tx.output.export or config.install_free_tensors
  1985. ):
  1986. # Export corner case - look at test_repros.py test_inlining_cornercase
  1987. new_source = UnspecializedBuiltinNNModuleSource(self.source)
  1988. result = UnspecializedBuiltinNNModuleVariable(value, source=new_source)
  1989. install_guard(new_source.make_guard(GuardBuilder.TYPE_MATCH))
  1990. else:
  1991. new_source = self.source
  1992. if config.inline_inbuilt_nn_modules and (
  1993. not self.tx.output.export or config.install_free_tensors
  1994. ):
  1995. # Export corner case - look at test_repros.py test_inlining_cornercase
  1996. new_source = UnspecializedNNModuleSource(self.source)
  1997. result = UnspecializedNNModuleVariable(value, source=new_source)
  1998. install_guard(new_source.make_guard(GuardBuilder.TYPE_MATCH))
  1999. self.tx.output.add_fqn_info_for_inlined_modules(value, self.source)
  2000. if not SideEffects.cls_supports_mutation_side_effects(type(value)):
  2001. # don't allow STORE_ATTR mutation with custom __setattr__
  2002. return result
  2003. return self.tx.output.side_effects.track_object_existing(value, result)
  2004. elif issubclass(
  2005. value.__class__, torch.nn.parallel.distributed.DistributedDataParallel
  2006. ):
  2007. self.install_guards(GuardBuilder.TYPE_MATCH)
  2008. return UnspecializedNNModuleVariable(value, source=self.get_source())
  2009. else:
  2010. return self.tx.output.register_attr_or_module(
  2011. value,
  2012. self.name,
  2013. source=self.get_source(),
  2014. # Guards are added inside register_attr_or_module
  2015. )
  2016. def wrap_literal(self, value: object) -> VariableTracker:
  2017. if type(value) is int:
  2018. assert isinstance(value, int)
  2019. # allowlist has higher precedence over specialization control.
  2020. if is_dynamic_source(self.source.name):
  2021. log.debug("%s marked dynamic via source whitelist", self.source.name)
  2022. return self.wrap_symint(value, dynamism=DimDynamic.DYNAMIC)
  2023. if is_unbacked_source(self.source.name):
  2024. log.debug("%s marked unbacked via source whitelist", self.source.name)
  2025. return self.wrap_symint(value, dynamism=DimDynamic.UNBACKED)
  2026. if not config.specialize_int:
  2027. # unspecializing int by default, but still
  2028. # specialize for the following conditions
  2029. if is_int_specialization_case(value, self.source):
  2030. recompile_hint = None
  2031. if (
  2032. self.source.guard_source.is_unspecialized_builtin_nn_module()
  2033. or self.source.guard_source.is_unspecialized_nn_module()
  2034. ):
  2035. # This means that it is an integer from a NN module.
  2036. # Dynamo considers nn module int attributes to be static
  2037. # (a good heuristic). But a user might want to mark the
  2038. # int attribute to be a symint, so track this integer
  2039. # for recompilation later.
  2040. recompile_hint = (
  2041. "torch.compile considers integer attributes of the nn.Module to be static. "
  2042. "If you are observing recompilation, you might want to make this integer dynamic "
  2043. "using torch._dynamo.config.allow_unspec_int_on_nn_module = True, or convert this "
  2044. "integer into a tensor."
  2045. )
  2046. process_automatic_dynamic(
  2047. self.tx,
  2048. self.source.name,
  2049. FrameStateSizeEntry.make_scalar(value),
  2050. is_unspecialized_nn_module=self.source.guard_source.is_unspecialized_nn_module(),
  2051. )
  2052. self.install_guards(
  2053. functools.partial(
  2054. GuardBuilder.EQUALS_MATCH, recompile_hint=recompile_hint
  2055. )
  2056. )
  2057. return ConstantVariable.create(value=value, source=self.source)
  2058. return self._wrap_lazy_constant(value, self._wrap_symint_for_lazy)
  2059. return self._wrap_lazy_constant(value)
  2060. elif type(value) is float:
  2061. assert isinstance(value, float)
  2062. if not config.specialize_float:
  2063. return self._wrap_lazy_constant(value, self._wrap_symfloat_for_lazy)
  2064. return self._wrap_lazy_constant(value)
  2065. elif type(value) in (bool, str):
  2066. assert isinstance(value, (bool, str))
  2067. return self._wrap_lazy_constant(value)
  2068. else:
  2069. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  2070. result = ConstantVariable.create(value=value, source=self.source)
  2071. if isinstance(value, (list, set)):
  2072. return self.tx.output.side_effects.track_mutable(value, result)
  2073. return result
  2074. def _wrap_symint_for_lazy(self, value: int) -> VariableTracker:
  2075. return self.wrap_symint(value)
  2076. def _wrap_symfloat_for_lazy(self, value: float) -> VariableTracker:
  2077. return self.wrap_symfloat(value)
  2078. @overload
  2079. def _wrap_lazy_constant(
  2080. self,
  2081. value: int,
  2082. wrap_fn: Callable[[int], VariableTracker],
  2083. ) -> VariableTracker: ...
  2084. @overload
  2085. def _wrap_lazy_constant(
  2086. self,
  2087. value: float,
  2088. wrap_fn: Callable[[float], VariableTracker],
  2089. ) -> VariableTracker: ...
  2090. @overload
  2091. def _wrap_lazy_constant(
  2092. self,
  2093. value: Union[int, float, bool, str],
  2094. wrap_fn: None = None,
  2095. ) -> VariableTracker: ...
  2096. def _wrap_lazy_constant(
  2097. self,
  2098. value: Union[int, float, bool, str],
  2099. wrap_fn: Optional[Callable[[Any], VariableTracker]] = None,
  2100. ) -> VariableTracker:
  2101. """Wrap a primitive constant, deferring guard installation if allowed."""
  2102. if not self.allow_lazy_constant:
  2103. if wrap_fn is not None:
  2104. return wrap_fn(value)
  2105. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  2106. return ConstantVariable.create(value=value, source=self.source)
  2107. return LazyConstantVariable.create(value, source=self.source)
  2108. def assert_not_wrapped_by_this_graph(self, value: torch.Tensor) -> None:
  2109. if is_fake(value) and maybe_get_fake_mode(value) is self.tx.fake_mode:
  2110. raise InternalTorchDynamoError(
  2111. "Cannot wrap a Tensor that has already been",
  2112. "wrapped by this instance of Dynamo",
  2113. )
  2114. def wrap_tensor(self, value: torch.Tensor) -> VariableTracker:
  2115. source = self.get_source()
  2116. # We cannot already be tracking the tensor, which implies
  2117. # it would have already been wrapped
  2118. assert value not in self.tx.output.side_effects
  2119. is_static_input = get_static_address_type(value) is not None
  2120. if (
  2121. config.inline_inbuilt_nn_modules
  2122. and not is_static_input
  2123. and (
  2124. isinstance(value, torch.nn.Parameter)
  2125. # mark tensor attributes of nn modules static. This is done to keep inline_inbuilt_nn_modules behavior
  2126. # compatible with previous behavior.
  2127. or (source and source.guard_source.is_unspecialized_nn_module())
  2128. )
  2129. ):
  2130. self.mark_static_input(value, guard=is_parameter_freezing())
  2131. is_static_input = True
  2132. # Install any tensors which are "free" variables; that is:
  2133. # 1. Globals
  2134. # 2. NonLocals
  2135. # 3. tensors that are attributes of nn module
  2136. should_install_free_tensor = config.install_free_tensors and (
  2137. is_from_global_source(source)
  2138. or is_from_nonlocal_source(source)
  2139. or is_from_unspecialized_nn_module_source(source)
  2140. )
  2141. make_graph_attribute = is_static_input and (
  2142. not config.inline_inbuilt_nn_modules
  2143. or is_parameter_freezing()
  2144. or torch._dynamo.config.prepare_freezing
  2145. )
  2146. if should_install_free_tensor or (
  2147. (source.guard_source.is_specialized_nn_module() or make_graph_attribute)
  2148. and not source.guard_source.is_fsdp_module()
  2149. ):
  2150. self.assert_not_wrapped_by_this_graph(value)
  2151. return self.tx.output.register_attr_or_module(
  2152. value, self.name, source=source
  2153. )
  2154. if get_static_address_type(value) == "guarded":
  2155. # If it's a guarded tensor, we can install the parameter directly
  2156. # into the Fx graph instead of lifting it as an input. Lifting
  2157. # offers no benefit, such as regional compilation, since we still
  2158. # guard on the tensor's ID. Moreover, installing it in the Fx graph
  2159. # eliminates the pre-graph bytecode required to extract the tensor
  2160. # from locals/globals, reducing overhead. This can lead to
  2161. # significant cost savings, especially for optimizers handling many
  2162. # tensors.
  2163. self.install_guards(GuardBuilder.ID_MATCH)
  2164. self.assert_not_wrapped_by_this_graph(value)
  2165. return self.tx.output.register_attr_or_module(
  2166. value, self.name, source=source
  2167. )
  2168. if is_constant_source(source):
  2169. self.assert_not_wrapped_by_this_graph(value)
  2170. return self.tx.output.register_attr_or_module(
  2171. value,
  2172. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  2173. source=source,
  2174. # Guards are added inside register_attr_or_module
  2175. )
  2176. # NB: this just says we accessed a tensor from the same source again
  2177. # (e.g., a tensor lives in a global foo, and we LOAD_GLOBAL it twice).
  2178. # This is distinct from two distinct sources mapping to the same
  2179. # Tensor (per id())! No guard is necessary here. See below for the
  2180. # other case.
  2181. is_duplicate_tensor = source in self.tx.output.input_source_to_var
  2182. if is_duplicate_tensor:
  2183. return self.tx.output.input_source_to_var[source]
  2184. options = {}
  2185. subclass_type = infer_subclass_type(value)
  2186. if subclass_type is not None:
  2187. self.install_guards(GuardBuilder.TYPE_MATCH)
  2188. if get_static_address_type(value) == "guarded":
  2189. self.install_guards(GuardBuilder.ID_MATCH)
  2190. # By this point, we should have deduplicated all tensors
  2191. self.assert_not_wrapped_by_this_graph(value)
  2192. if (
  2193. isinstance(value, torch.Tensor)
  2194. and value.is_nested
  2195. and not isinstance(value, torch.nested._internal.nested_tensor.NestedTensor)
  2196. ):
  2197. unimplemented(
  2198. gb_type="Attempted to wrap strided NestedTensor",
  2199. context="",
  2200. explanation="torch.compile does not support strided NestedTensor",
  2201. hints=[],
  2202. )
  2203. # TODO(pearu,sparse-team) - Add the corresponding SPARSE_TENSOR_MATCH guards
  2204. if (
  2205. isinstance(value, torch.Tensor)
  2206. and is_sparse_any(value)
  2207. and (not self.tx.export or not config.capture_sparse_compute)
  2208. ):
  2209. # A hot fix for sparse tensors + torch.compile. Support for
  2210. # export + sparsity is being added but we need to create
  2211. # SPARSE_TENSOR_GUARDS for guards to work properly.
  2212. unimplemented(
  2213. gb_type="Attempted to wrap sparse Tensor",
  2214. context="",
  2215. explanation="torch.compile does not support sparse Tensors",
  2216. hints=[*graph_break_hints.SPARSE_TENSOR],
  2217. )
  2218. if (
  2219. safe_has_grad(value)
  2220. and safe_grad(value) is not None
  2221. # type: ignore[attr-defined]
  2222. and value.dtype != safe_grad(value).dtype
  2223. ):
  2224. safe_grad_val = safe_grad(value)
  2225. grad_str = str(safe_grad_val.dtype) if safe_grad_val is not None else "None"
  2226. unimplemented(
  2227. gb_type="dtype mismatch between tensor and its gradient",
  2228. context=f"tensor dtype: {value.dtype}; grad dtype: {grad_str}",
  2229. explanation="Inconsistent dtype between tensor and its gradient. "
  2230. "This can happen in FSDP and crashes meta tensor creation.",
  2231. hints=[*graph_break_hints.SUPPORTABLE],
  2232. )
  2233. # tx.output has multiple tracers if we're introspecting HigherOrderOperator.
  2234. # When we've discovered an untracked tensor, then we actually need
  2235. # to get Dynamo to track the tensor (which is what this function does)
  2236. # and put it as a graph input on the root tracer. Later on,
  2237. # if the input is actually used in the body of the HigherOrderOperator,
  2238. # then the relevant SubgraphTracer will lift it to being an input of
  2239. # the subgraph.
  2240. # See NOTE [HigherOrderOperator tracing design] for more details.
  2241. example_value = wrap_to_fake_tensor_and_record(
  2242. value, tx=self.tx, is_tensor=True, source=source
  2243. )
  2244. tensor_proxy = self.tx.output.root_tracer.create_graph_input(
  2245. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  2246. type(value),
  2247. example_value,
  2248. source=source,
  2249. )
  2250. cache_real_value_when_export(self.tx, tensor_proxy, value)
  2251. tensor_variable = wrap_fx_proxy(
  2252. tx=self.tx,
  2253. proxy=tensor_proxy,
  2254. example_value=example_value,
  2255. subclass_type=subclass_type,
  2256. source=source,
  2257. **options,
  2258. )
  2259. if value._is_view():
  2260. # If value is a view, add its base tensor to the tracked fakes list.
  2261. # This is so we are able to access the correct source for its symbolic
  2262. # shape values, in case we need them.
  2263. wrap_to_fake_tensor_and_record(
  2264. value._base,
  2265. tx=self.tx,
  2266. source=AttrSource(source, "_base"),
  2267. is_tensor=True,
  2268. )
  2269. guard_type = GuardBuilder.TENSOR_MATCH
  2270. if isinstance(source, GradSource) and is_from_optimizer_source(source):
  2271. guard_type = GuardBuilder.NOT_NONE_MATCH
  2272. is_dtensor = torch.distributed.is_available() and isinstance(
  2273. value, torch.distributed.tensor.DTensor
  2274. )
  2275. if not is_dtensor:
  2276. # We guard on the _local_tensor and the _spec, and therefore we dont
  2277. # have to guard on the outer DTensor.
  2278. self.install_guards(
  2279. functools.partial(
  2280. guard_type,
  2281. value=(
  2282. value
  2283. if isinstance(source, NumpyTensorSource)
  2284. else TensorWeakRef(value)
  2285. ),
  2286. )
  2287. )
  2288. # We install TYPE_MATCH guards for traceable wrapper subclass object,
  2289. # and recursively install corresponding guard for each inner attribute.
  2290. if is_traceable_wrapper_subclass(value):
  2291. # Tensor subclass guards are very expensive because they are
  2292. # implemented in Python. Since DTensor is PyTorch-maintained class,
  2293. # we can skip a lot of these guards.
  2294. if is_dtensor:
  2295. self.install_guards(GuardBuilder.TYPE_MATCH)
  2296. # The inner tensor name is always _local_tensor. If its not, we
  2297. # raise assertion to update the check accordingly.
  2298. inner_tensor_name = value.__tensor_flatten__()[0][0]
  2299. if inner_tensor_name != "_local_tensor":
  2300. raise RuntimeError(
  2301. "Expecting Dtensor inner tensor name to be _local_tensor"
  2302. )
  2303. # Now selectively guard on the flattening context
  2304. flattening_ctx = value.__tensor_flatten__()[1]
  2305. # This is supposed to be (self._spec, self.requires_grad)
  2306. if not (
  2307. len(flattening_ctx) == 2
  2308. and flattening_ctx[0] == value._spec
  2309. and flattening_ctx[1] == value.requires_grad
  2310. ):
  2311. # If not, raise an assertion to update to the new guards
  2312. raise RuntimeError(
  2313. "Expecting Dtensor flattening ctx to be _spec, requires_grad"
  2314. )
  2315. # Guard on the dtensor spec
  2316. install_guard(
  2317. AttrSource(self.source, "_spec").make_guard(
  2318. GuardBuilder.DTENSOR_SPEC_MATCH
  2319. )
  2320. )
  2321. # Move this to C++
  2322. install_guard(
  2323. AttrSource(self.source, "requires_grad").make_guard(
  2324. GuardBuilder.EQUALS_MATCH
  2325. )
  2326. )
  2327. else:
  2328. self.install_guards(GuardBuilder.TENSOR_SUBCLASS_METADATA_MATCH)
  2329. self.install_guards(GuardBuilder.TYPE_MATCH)
  2330. install_guard(
  2331. SubclassAttrListSource(source).make_guard(GuardBuilder.EQUALS_MATCH)
  2332. )
  2333. attrs, _ = value.__tensor_flatten__()
  2334. for attr in attrs:
  2335. inner_value = getattr(value, attr)
  2336. inner_source = AttrSource(self.source, attr)
  2337. LazyVariableTracker.realize_all(
  2338. VariableBuilder(self.tx, inner_source)(inner_value)
  2339. )
  2340. self.tx.output.input_source_to_var[source] = tensor_variable
  2341. assert "tensor_dict" not in tensor_proxy.node.meta
  2342. tensor_proxy.node.meta["tensor_dict"] = _extract_tensor_dict(value)
  2343. # Note: this information is conveyed via subclass_type now
  2344. # type: ignore[attr-defined]
  2345. fake_tensor_value = tensor_variable.proxy.node.meta["example_value"]
  2346. if maybe_get_fake_mode(fake_tensor_value) is not self.tx.fake_mode:
  2347. raise InternalTorchDynamoError("Wrapped Tensor must be this graph's fake")
  2348. grapharg = GraphArg(source, value, False, fake_tensor_value)
  2349. tensor_proxy.node.meta["grapharg"] = grapharg
  2350. return tensor_variable
  2351. def wrap_numpy_ndarray(self, value: Any) -> VariableTracker:
  2352. assert np is not None
  2353. assert isinstance(value, np.ndarray)
  2354. source = NumpyTensorSource(self.get_source())
  2355. from torch._numpy import _util
  2356. readonly = not value.flags.writeable
  2357. if readonly:
  2358. try:
  2359. value.flags.writeable = True
  2360. except ValueError:
  2361. # One can not easily make nditer elements writable,
  2362. # but warning is not the end of the world
  2363. assert isinstance(value.base, np.nditer)
  2364. tensor_value = None
  2365. with torch_function_mode_stack_state_mgr.temp_restore_stack():
  2366. try:
  2367. tensor_value = _util._try_convert_to_tensor(value)
  2368. if readonly:
  2369. from torch._prims_common import clone_preserve_strides
  2370. tensor_value = clone_preserve_strides(tensor_value)
  2371. except NotImplementedError as e:
  2372. # failed to convert to tensor, graph break
  2373. unimplemented(
  2374. gb_type="failed to convert numpy.ndarray to Tensor",
  2375. context=str(value),
  2376. explanation="Exception encountered when attempting to convert numpy.ndarray to Tensor",
  2377. hints=[],
  2378. from_exc=e,
  2379. )
  2380. assert tensor_value is not None
  2381. # We do this because we want the full behavior of guarding the numpy ndarray as if it were
  2382. # a tensor. It's a little annoying to make a VT to throw out, but there's so many side effects here
  2383. # that there's not another great way to do this atm.
  2384. # This creates the right graphargs, as well as registration for guards in tensor names and shape env.
  2385. LazyVariableTracker.realize_all(VariableBuilder(self.tx, source)(tensor_value))
  2386. example_value = wrap_to_fake_tensor_and_record(
  2387. tensor_value,
  2388. tx=self.tx,
  2389. is_tensor=False,
  2390. source=source,
  2391. )
  2392. proxy = self.tx.output.root_tracer.create_graph_input(
  2393. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  2394. type(tensor_value),
  2395. example_value,
  2396. source=source,
  2397. )
  2398. cache_real_value_when_export(self.tx, proxy, tensor_value)
  2399. options = {"source": source}
  2400. numpy_ndarray_variable = wrap_fx_proxy_cls(
  2401. target_cls=NumpyNdarrayVariable,
  2402. tx=self.tx,
  2403. proxy=proxy,
  2404. example_value=example_value,
  2405. subclass_type=None,
  2406. **options,
  2407. )
  2408. self.tx.output.input_source_to_var[source] = numpy_ndarray_variable
  2409. # type: ignore[attr-defined]
  2410. example_value = numpy_ndarray_variable.proxy.node.meta["example_value"]
  2411. # pass_arg_as_tensor should be true because we are wrapping a np.ndarray as argument input, and it needs to be
  2412. # converted to a tensor.
  2413. grapharg = GraphArg(
  2414. source,
  2415. tensor_value,
  2416. pass_arg_as_tensor=True,
  2417. fake_tensor=example_value,
  2418. is_tensor=True,
  2419. example_strong_ref=tensor_value,
  2420. )
  2421. proxy.node.meta["grapharg"] = grapharg
  2422. # TODO - Why do we need to set the source of the np ndarray vt back to
  2423. # original source. Many tests fails.
  2424. numpy_ndarray_variable.source = self.source
  2425. return numpy_ndarray_variable
  2426. def wrap_symint(
  2427. self,
  2428. value: int,
  2429. dynamism: DimDynamic | None = None,
  2430. context: SymIntSymbolicContext | None = None,
  2431. ) -> VariableTracker:
  2432. assert type(value) is int
  2433. if self.name in self.tx.output.unspec_variable_map:
  2434. return self.tx.output.unspec_variable_map[self.name]
  2435. shape_env = self.tx.output.shape_env
  2436. if TracingContext.get().force_unspec_int_unbacked_size_like:
  2437. wrapped_value = shape_env.create_unbacked_symint()
  2438. _constrain_range_for_size(wrapped_value)
  2439. self.tx.output.tracked_fakes.append(
  2440. TrackedFake(wrapped_value, self.source, None)
  2441. )
  2442. # NB: We do not do float. For motivation, see
  2443. # https://docs.google.com/document/d/1INSCdYu1PxXcr43HrD82OudeEuS-qxQe1yZmLg2wy6A/edit
  2444. # but the general idea is that we generate kernels that can
  2445. # take unspecialized floats and use them in sizevar computation
  2446. elif not is_constant_source(self.get_source()):
  2447. if dynamism is None and torch._dynamo.config.specialize_int:
  2448. # If specialize_int is False, also return
  2449. # a constant (but this should have been handled
  2450. # in the caller, TBH). But if `dynamism` is set, then actually
  2451. # turn it into a symint
  2452. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  2453. return ConstantVariable.create(value=value, source=self.source)
  2454. name = self.source.name
  2455. frame_state_entry = process_automatic_dynamic(
  2456. self.tx,
  2457. name,
  2458. FrameStateSizeEntry.make_scalar(value),
  2459. is_unspecialized_nn_module=self.source.guard_source.is_unspecialized_nn_module(),
  2460. )
  2461. # TODO: This should be dynamic, as we in general do not
  2462. # know if bare integers are actually going to be sizevars
  2463. # and it is inappropriate to eagerly duck size them with
  2464. # real sizevars
  2465. normalized_source_name = normalize_source_name(self.source.name)
  2466. base_source = self.source
  2467. if isinstance(base_source, ChainedSource):
  2468. base_source = base_source.get_base()
  2469. if dynamism is not None:
  2470. dynamic_dim = dynamism
  2471. elif (
  2472. config.automatic_dynamic_shapes
  2473. and frame_state_entry.scalar is auto_dynamic
  2474. ):
  2475. set_feature_use("dynamo.automatic_dynamic_shapes", True)
  2476. dynamic_dim = get_automatic_dynamic_shapes_mark_as()
  2477. elif (
  2478. isinstance(base_source, LocalSource)
  2479. and base_source.dynamism is not None
  2480. # pyrefly: ignore[no-matching-overload]
  2481. and dict(base_source.dynamism).get(normalized_source_name, {0: False})[
  2482. 0
  2483. ]
  2484. ) or not config.assume_static_by_default:
  2485. dynamic_dim = DimDynamic.DYNAMIC
  2486. else: # assume_static_by_default
  2487. # TODO: dynamic_dim = DimDynamic.STATIC should work but
  2488. # for some reason it doesn't
  2489. if frame_state_entry.scalar is auto_dynamic:
  2490. set_feature_use("dynamo.automatic_dynamic_shapes", False)
  2491. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  2492. return ConstantVariable.create(value=value)
  2493. wrapped_value = shape_env.create_unspecified_symint_and_symbol(
  2494. value,
  2495. source=self.source,
  2496. dynamic_dim=dynamic_dim,
  2497. )
  2498. self.tx.output.tracked_fakes.append(
  2499. TrackedFake(wrapped_value, self.source, context)
  2500. )
  2501. else:
  2502. assert is_constant_source(self.get_source())
  2503. # TODO: Do I actually need guard for constant source?
  2504. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  2505. return ConstantVariable.create(value=value, source=self.source)
  2506. assert not isinstance(self.get_source(), RandomValueSource)
  2507. install_guard(self.get_source().make_guard(GuardBuilder.TYPE_MATCH))
  2508. options = {"source": self.get_source()}
  2509. proxy = self.tx.output.root_tracer.create_graph_input(
  2510. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  2511. type(wrapped_value),
  2512. wrapped_value,
  2513. source=self.get_source(),
  2514. )
  2515. sym_expr = wrapped_value.node.expr
  2516. assert isinstance(sym_expr, sympy.Symbol), f"{sym_expr} is not a basic Symbol."
  2517. self.tx.output.root_tracer.bound_symbols[sym_expr] = proxy
  2518. unspec_var = SymNodeVariable.create(self.tx, proxy, wrapped_value, **options)
  2519. # type: ignore[assignment]
  2520. self.tx.output.unspec_variable_map[self.name] = unspec_var
  2521. if not is_constant_source(self.get_source()):
  2522. proxy.node.meta["grapharg"] = GraphArg(
  2523. self.get_source(),
  2524. wrapped_value,
  2525. pass_arg_as_tensor=False,
  2526. fake_tensor=None,
  2527. is_tensor=False,
  2528. example_strong_ref=wrapped_value,
  2529. )
  2530. return unspec_var
  2531. def wrap_symfloat(self, value: float) -> VariableTracker:
  2532. # SymFloat wrapping is special. We first wrap it in the same way we
  2533. # do an unspecialized primitive, and then we item() it into a
  2534. # SymFloat. Removal of the item() call is left to a later FX pass,
  2535. # mostly because that pass is more easily done after we have lowered
  2536. # to ATen ops. (Dynamo doesn't do decomposition right now).
  2537. if self.name in self.tx.output.unspec_variable_map:
  2538. return self.tx.output.unspec_variable_map[self.name]
  2539. frame_state_entry = process_automatic_dynamic(
  2540. self.tx,
  2541. self.source.name,
  2542. # type: ignore[arg-type]
  2543. FrameStateSizeEntry.make_scalar(value),
  2544. is_unspecialized_nn_module=self.source.guard_source.is_unspecialized_nn_module(),
  2545. )
  2546. # NB: we specialize on nan input, because our guard modeling in
  2547. # ShapeEnv cannot deal with nan
  2548. if (
  2549. torch._dynamo.config.specialize_float
  2550. or is_constant_source(self.get_source())
  2551. or math.isnan(value)
  2552. or math.isinf(value)
  2553. # We don't support cudagraphs for now. Without this cudagraphs
  2554. # break because they expect all cuda inputs but our tensorified
  2555. # float will be a f64[] cpu tensor. Fixes the following test
  2556. # when specialize_float=False
  2557. # python test/inductor/test_compiled_optimizers.py CompiledOptimizerTests.test_rmsprop_weight_decay_maximize_capturable_cuda # noqa: B950
  2558. or torch._inductor.config.triton.cudagraphs
  2559. or justknobs_check("pytorch/compiler:unspecialize_float_killswitch", False)
  2560. or (
  2561. config.assume_static_by_default
  2562. and frame_state_entry.scalar is not auto_dynamic
  2563. )
  2564. ):
  2565. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  2566. return ConstantVariable.create(value=value, source=self.source)
  2567. # NB: At the point we've gotten here, we don't assume static by
  2568. # default. Since we have a guard mechanism, there isn't really any
  2569. # downside to trying to be dynamic for float all the time. Unlike
  2570. # ints, this won't make codegen perf worse. Modest cost to compile
  2571. # time.
  2572. wrapped_value = torch.tensor(value, dtype=torch.float64)
  2573. # We don't support specializing floats for grad checking tensors
  2574. # See https://github.com/pytorch/pytorch/pull/140828 for more
  2575. # context.
  2576. if torch._C._functorch.is_gradtrackingtensor(wrapped_value):
  2577. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  2578. return ConstantVariable.create(value=value, source=self.source)
  2579. # TODO: Switch RandomValueSource over to use this, this is more
  2580. # accurate
  2581. assert not isinstance(self.get_source(), RandomValueSource)
  2582. install_guard(self.get_source().make_guard(GuardBuilder.TYPE_MATCH))
  2583. # The FloatTensorSource here is just for pedantic correctness: if you
  2584. # guard against an UnspecializedPythonVariable, you need to guard
  2585. # against the tensor-ified version of the local, otherwise it's not a
  2586. # Tensor. However, we never let the UnspecializedPythonVariable escape
  2587. # here, so there should never actually be any guards against this
  2588. # source.
  2589. source = FloatTensorSource(self.get_source())
  2590. options = {"source": source, "raw_value": value}
  2591. # TODO: Maybe the tensor-ification should be built into the source,
  2592. # rather than by special pattern match
  2593. example_value = wrap_to_fake_tensor_and_record(
  2594. wrapped_value, tx=self.tx, is_tensor=False, source=source
  2595. )
  2596. proxy = self.tx.output.root_tracer.create_graph_input(
  2597. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  2598. type(wrapped_value),
  2599. example_value,
  2600. source=source,
  2601. )
  2602. cache_real_value_when_export(self.tx, proxy, wrapped_value)
  2603. unspec_var = wrap_fx_proxy_cls(
  2604. UnspecializedPythonVariable,
  2605. tx=self.tx,
  2606. proxy=proxy,
  2607. example_value=example_value,
  2608. subclass_type=None,
  2609. **options,
  2610. )
  2611. assert isinstance(unspec_var, UnspecializedPythonVariable)
  2612. self.tx.output.unspec_variable_map[self.name] = unspec_var
  2613. if self.tx.export and not isinstance(self.get_source(), LocalSource):
  2614. raise AssertionError(
  2615. f"Dynamo attempts to add additional input during export: value={wrapped_value}, source={self.get_source()}"
  2616. )
  2617. fake_tensor_value = None
  2618. example_value = unspec_var.proxy.node.meta["example_value"]
  2619. assert is_fake(example_value)
  2620. fake_tensor_value = example_value
  2621. # type: ignore[attr-defined]
  2622. assert fake_tensor_value.fake_mode is self.tx.fake_mode, (
  2623. f"fake mode ({fake_tensor_value.fake_mode}) from fake tensor metadata doesn't match mode"
  2624. "({self.tx.fake_mode}) from InstructionTranslator"
  2625. )
  2626. # There's something a bit incoherent about pass_arg_as_tensor,
  2627. # specifically regarding sources.
  2628. #
  2629. # Specifically, suppose we have "x: float" local argument. We
  2630. # eventually end up with an UnspecializedPythonVariable denoting
  2631. # torch.as_tensor(x)... but it's source is still L['x'] (which if you
  2632. # accessed it directly is a float!) So you gotta be careful when
  2633. # setting up your guards, because it's still going to be a float at
  2634. # this point, the conversion happens only precisely at the point we're
  2635. # actually calling the FX graph. This happens to be what we want for
  2636. # shape guard generation, but it's kind of unintuitive.
  2637. proxy.node.meta["grapharg"] = GraphArg(
  2638. self.get_source(),
  2639. wrapped_value,
  2640. pass_arg_as_tensor=True,
  2641. # type: ignore[arg-type]
  2642. fake_tensor=fake_tensor_value,
  2643. is_tensor=False,
  2644. example_strong_ref=wrapped_value,
  2645. )
  2646. # Directly do item to bypass capture_scalar_outputs
  2647. r = wrap_fx_proxy(
  2648. self.tx,
  2649. self.tx.output.create_proxy(
  2650. "call_method",
  2651. "item",
  2652. *proxy_args_kwargs([unspec_var], {}),
  2653. ),
  2654. )
  2655. # type: ignore[attr-defined]
  2656. self.tx.output.tracked_fakes.append(TrackedFake(r.sym_num, self.source, None))
  2657. get_metrics_context().set("tensorify_float_attempt", True, overwrite=True)
  2658. return r
  2659. def wrap_unspecialized_primitive(self, value: Any) -> VariableTracker:
  2660. if self.name in self.tx.output.unspec_variable_map:
  2661. return self.tx.output.unspec_variable_map[self.name]
  2662. wrapped_value = torch.tensor(value)
  2663. if not isinstance(self.get_source(), RandomValueSource):
  2664. install_guard(self.get_source().make_guard(GuardBuilder.TYPE_MATCH))
  2665. options = {"source": self.get_source()}
  2666. options.update({"raw_value": value})
  2667. example_value = wrap_to_fake_tensor_and_record(
  2668. wrapped_value, tx=self.tx, is_tensor=False, source=self.get_source()
  2669. )
  2670. proxy = self.tx.output.root_tracer.create_graph_input(
  2671. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  2672. type(wrapped_value),
  2673. example_value,
  2674. source=self.get_source(),
  2675. )
  2676. cache_real_value_when_export(self.tx, proxy, wrapped_value)
  2677. unspec_var = wrap_fx_proxy_cls(
  2678. UnspecializedPythonVariable,
  2679. tx=self.tx,
  2680. proxy=proxy,
  2681. example_value=example_value,
  2682. subclass_type=None,
  2683. **options,
  2684. )
  2685. # type: ignore[assignment]
  2686. self.tx.output.unspec_variable_map[self.name] = unspec_var
  2687. if not is_constant_source(self.get_source()):
  2688. if self.tx.export and not isinstance(self.get_source(), LocalSource):
  2689. raise AssertionError(
  2690. f"Dynamo attempts to add additional input during export: value={wrapped_value}, source={self.get_source()}"
  2691. )
  2692. fake_tensor_value = None
  2693. if unspec_var.is_python_constant():
  2694. # TODO: when can this happen?
  2695. example_value = unspec_var.as_python_constant()
  2696. else:
  2697. # type: ignore[attr-defined]
  2698. example_value = unspec_var.proxy.node.meta["example_value"]
  2699. assert is_fake(example_value)
  2700. fake_tensor_value = example_value
  2701. # type: ignore[attr-defined]
  2702. assert fake_tensor_value.fake_mode is self.tx.fake_mode, (
  2703. f"fake mode ({fake_tensor_value.fake_mode}) from fake tensor metadata doesn't match mode"
  2704. "({self.tx.fake_mode}) from InstructionTranslator"
  2705. )
  2706. proxy.node.meta["grapharg"] = GraphArg(
  2707. self.get_source(),
  2708. wrapped_value,
  2709. pass_arg_as_tensor=True,
  2710. # type: ignore[arg-type]
  2711. fake_tensor=fake_tensor_value,
  2712. is_tensor=False,
  2713. example_strong_ref=wrapped_value,
  2714. )
  2715. return unspec_var
  2716. def _dataclasses_fields_lambda(obj: VariableTracker) -> TupleVariable:
  2717. value = None
  2718. if isinstance(obj, UserDefinedObjectVariable):
  2719. value = obj.value
  2720. else:
  2721. unimplemented(
  2722. gb_type="dataclass fields failure",
  2723. context=f"obj: {obj}; variable type: {type(obj)}",
  2724. explanation=f"Dataclass fields handling fails for {obj}. Expected it to be a user-defined object.",
  2725. hints=[],
  2726. )
  2727. assert value is not None
  2728. items = []
  2729. # type: ignore[arg-type]
  2730. for field in dataclasses.fields(value):
  2731. source = None
  2732. if obj.source:
  2733. base_src = AttrSource(obj.source, "__dataclass_fields__")
  2734. source = DictGetItemSource(base_src, field.name)
  2735. items.append(UserDefinedObjectVariable(field, source=source))
  2736. # pyrefly: ignore [bad-argument-type]
  2737. return TupleVariable(items)
  2738. def _clone_input(value: Any, fake_mode: FakeTensorMode | None) -> Any:
  2739. if isinstance(value, torch.Tensor):
  2740. # tensor subclasses will not be converted to FakeTensors and need to be cloned
  2741. if not (
  2742. isinstance(value, FakeTensor)
  2743. or (
  2744. # Is functional tensor fakeified by this instance of Dynamo
  2745. torch._is_functional_tensor(value)
  2746. and maybe_get_fake_mode(value) is fake_mode
  2747. )
  2748. or value.is_nested
  2749. ):
  2750. # NB: ensure strides are preserved
  2751. value = clone_input(value)
  2752. return value
  2753. def wrap_fx_proxy(
  2754. tx: "InstructionTranslatorBase",
  2755. proxy: Any,
  2756. example_value: Any | None = None,
  2757. subclass_type: type | None = None,
  2758. **options: Any,
  2759. ) -> VariableTracker:
  2760. kwargs = {
  2761. "tx": tx,
  2762. "proxy": proxy,
  2763. "example_value": example_value,
  2764. "subclass_type": subclass_type,
  2765. **options,
  2766. }
  2767. if subclass_type is None:
  2768. # pyrefly: ignore[bad-argument-type]
  2769. return wrap_fx_proxy_cls(target_cls=TensorVariable, **kwargs)
  2770. else:
  2771. # pyrefly: ignore[bad-argument-type]
  2772. result = wrap_fx_proxy_cls(target_cls=TensorWithTFOverrideVariable, **kwargs)
  2773. # type: ignore[attr-defined]
  2774. result.install_global(tx)
  2775. return result
  2776. def cache_real_value_when_export(
  2777. tx: "InstructionTranslatorBase", proxy: Any, example_value: Any
  2778. ) -> None:
  2779. if tx.export:
  2780. # The legacy behavior for real value cache with subclasses was
  2781. # to perform a clone WITHOUT preserving the subclass. It's
  2782. # not entirely clear this is what you actually want though.
  2783. with torch._C.DisableTorchFunctionSubclass():
  2784. proxy.tracer.real_value_cache[proxy.node] = _clone_input(
  2785. example_value, tx.fake_mode
  2786. )
  2787. # Note: Unfortunate split due to some gross classes existing that subclass TensorVariable
  2788. # Should be compositional instead
  2789. #
  2790. # This is a horribly complicated function that does too many things, to
  2791. # explain what it does, let's first talk about the classic usage wrap_fx_proxy
  2792. # for a TensorVariable. There are two primary modes of use:
  2793. #
  2794. # 1. Wrapping a pre-existing Tensor. In this case, example_value is set
  2795. # to the pre-existing Tensor. (Note that this example_value will NOT
  2796. # be the final example_value we put into node.meta['example_value'],
  2797. # instead it is converted into a fake tensor using
  2798. # wrap_to_fake_tensor_and_record and registered as a graph input.)
  2799. #
  2800. # 2. "Wrapping" the result of some Tensor operation Dynamo traced over. In
  2801. # this case, example_value is None (and we are going to figure it out
  2802. # ourselves using FakeTensors, via get_fake_value, which will run
  2803. # the operation represented by the (singular!) FX node referenced by
  2804. # the passed in proxy.)
  2805. #
  2806. # The expectation is you end up with a Tensor output, and everything is
  2807. # straightforwardly traced into the graph.
  2808. #
  2809. # In all cases, the returned `TensorVariable` subclass will have an `example_value`
  2810. # and that `example_value` must be a `FakeTensor` produced by the currently running
  2811. # instance of Dynamo.
  2812. #
  2813. # Upon closer inspection, you may notice that there are a slurry of non-Tensor
  2814. # output cases in handle_traced_output. What gives? Well, we sometimes trace operations into the
  2815. # graph that don't involve tensors.
  2816. #
  2817. # * Some operators return tuples; we need to recursively handle their
  2818. # contents
  2819. #
  2820. # * Some operators have side effects that will affect subsequent AOTAutograd
  2821. # tracing but don't otherwise return anything.
  2822. #
  2823. # * Some operators return symbolic ints/floats/bools which can go in the
  2824. # graph and be traced (but only if they're actually symbolic! If they're
  2825. # static you don't want to put them in the graph, which means you
  2826. # shouldn't call this function.)
  2827. #
  2828. # The common theme is that you only use this function WHEN YOU ARE TRACING
  2829. # SOMETHING INTO THE GRAPH. This is sort of obvious, because you can't call
  2830. # this function without a proxy.
  2831. def wrap_fx_proxy_cls(
  2832. target_cls: type[VTTypeAlias],
  2833. tx: "InstructionTranslatorBase",
  2834. proxy: Any,
  2835. example_value: Any | None = None,
  2836. subclass_type: type | None = None,
  2837. **options: Any,
  2838. ) -> VTTypeAlias:
  2839. if example_value is None:
  2840. out: VTTypeAlias = _wrap_fx_proxy(
  2841. target_cls, tx, proxy, example_value, subclass_type, **options
  2842. )
  2843. elif isinstance(example_value, torch.Tensor):
  2844. out = _wrap_fx_preexisting_tensor(
  2845. target_cls, tx, proxy, example_value, subclass_type, **options
  2846. )
  2847. else:
  2848. # This will skip tracing an op and recursively reinvoke wrap_fx_proxy_cls on supported
  2849. # data structures. In essence this just handles tracing some other value which may
  2850. # contain Fake Tensors or is otherwise proxyable.
  2851. # pyrefly: ignore[bad-assignment]
  2852. out = handle_traced_output(
  2853. example_value, tx, proxy, options, subclass_type, target_cls
  2854. )
  2855. if (
  2856. isinstance(
  2857. out,
  2858. (
  2859. torch._dynamo.variables.TensorVariable,
  2860. torch._dynamo.variables.SymNodeVariable,
  2861. ),
  2862. )
  2863. and proxy.node.op != "placeholder"
  2864. ):
  2865. tx.output.current_tracer.record_tensor_or_symint_vt(out)
  2866. return out
  2867. # This is 1 above (wrapping a preexisting tensor)
  2868. def _wrap_fx_preexisting_tensor(
  2869. target_cls: type[VTTypeAlias],
  2870. tx: "InstructionTranslatorBase",
  2871. proxy: torch.fx.Proxy,
  2872. tensor: torch.Tensor,
  2873. subclass_type: type | None = None,
  2874. **options: Any,
  2875. ) -> VTTypeAlias:
  2876. from ..symbolic_convert import InstructionTranslatorBase
  2877. assert isinstance(tensor, torch.Tensor), (
  2878. f"_wrap_fx_preexisting_tensor expected tensor, got {type(tensor)}"
  2879. )
  2880. assert isinstance(tx, InstructionTranslatorBase)
  2881. if "guards" in options and options["guards"] is not None:
  2882. tx.output.guards.update(options["guards"])
  2883. # Placeholders always carry example_value in node.meta.
  2884. # non-placeholders always have no example_value in node.meta
  2885. if proxy.node.op == "placeholder":
  2886. assert "example_value" in proxy.node.meta, (
  2887. f"placeholder {proxy} doesn't have 'example_value' in node.meta"
  2888. )
  2889. else:
  2890. assert "example_value" not in proxy.node.meta, (
  2891. f"{proxy.node.meta['example_value']}"
  2892. )
  2893. # See NOTE: [Deferring tensor pack/unpack hooks until runtime]
  2894. with torch._dynamo.utils._disable_saved_tensors_hooks_during_tracing():
  2895. # Handle recursive calls here
  2896. if maybe_get_fake_mode(tensor) is tx.fake_mode:
  2897. pass
  2898. else:
  2899. cache_real_value_when_export(tx, proxy, tensor)
  2900. if tx.export:
  2901. # The legacy behavior for real value cache with subclasses was
  2902. # to perform a clone WITHOUT preserving the subclass. It's
  2903. # not entirely clear this is what you actually want though.
  2904. with torch._C.DisableTorchFunctionSubclass():
  2905. # type: ignore[attr-defined]
  2906. proxy.tracer.real_value_cache[proxy.node] = _clone_input(
  2907. tensor, tx.fake_mode
  2908. )
  2909. # NB: If we're ignoring subclass, then the expectation is you will
  2910. # take the returned TensorVariable and wrap it into a more
  2911. # accurate TensorVariable that is able to track subclass-ness;
  2912. # otherwise this is wrong!
  2913. kwargs = {
  2914. "is_tensor": target_cls
  2915. in (TensorVariable, TensorWithTFOverrideVariable),
  2916. }
  2917. assert "source" in options and options["source"] is not None
  2918. kwargs["source"] = options["source"]
  2919. # pyrefly: ignore[missing-argument, bad-argument-type]
  2920. tensor = wrap_to_fake_tensor_and_record(tensor, tx=tx, **kwargs)
  2921. if tensor.device.type != "meta" and (
  2922. maybe_get_fake_mode(tensor) is not tx.fake_mode
  2923. ):
  2924. raise InternalTorchDynamoError(
  2925. "`tensor` needs to be a `FakeTensor`"
  2926. f"wrapped by this instance of Dynamo. Found: {tensor}"
  2927. )
  2928. return construct_tensor_variable(
  2929. target_cls, tx, proxy, tensor, subclass_type, options
  2930. )
  2931. # This is 2 in the above comment (wrapping the output of a traced op)
  2932. def _wrap_fx_proxy(
  2933. target_cls: type[VTTypeAlias],
  2934. tx: "InstructionTranslatorBase",
  2935. proxy: torch.fx.Proxy,
  2936. example_value: Any | None = None,
  2937. subclass_type: type | None = None,
  2938. **options: Any,
  2939. ) -> VTTypeAlias:
  2940. from ..symbolic_convert import InstructionTranslatorBase
  2941. assert isinstance(tx, InstructionTranslatorBase)
  2942. if "guards" in options and options["guards"] is not None:
  2943. tx.output.guards.update(options["guards"])
  2944. assert "example_value" not in proxy.node.meta, f"{proxy.node.meta['example_value']}"
  2945. # See NOTE: [Deferring tensor pack/unpack hooks until runtime]
  2946. with torch._dynamo.utils._disable_saved_tensors_hooks_during_tracing():
  2947. # with preserve_rng_state():
  2948. # only allow_non_graph_fake in this instance because we handle the non-fake
  2949. # cases properly below.
  2950. example_value = get_fake_value(proxy.node, tx, allow_non_graph_fake=True)
  2951. # pyrefly: ignore[bad-return]
  2952. return handle_traced_output(
  2953. # type: ignore[arg-type]
  2954. example_value,
  2955. tx,
  2956. proxy,
  2957. options,
  2958. subclass_type,
  2959. target_cls,
  2960. )
  2961. # This handles wrapping of the output of an op traced into the graphs
  2962. def handle_traced_output(
  2963. example_value: Any,
  2964. tx: "InstructionTranslatorBase",
  2965. proxy: torch.fx.Proxy,
  2966. options: dict[str, Any],
  2967. subclass_type: type | None,
  2968. target_cls: type[VTTypeAlias],
  2969. ) -> VariableTracker:
  2970. import torch._functorch.vmap
  2971. import torch._subclasses.fake_tensor
  2972. import torch._utils
  2973. if isinstance(example_value, torch.Tensor):
  2974. # Check if the result is a sparse tensor -
  2975. # We generally don't support sparse tensor so better to graph break here
  2976. if is_sparse_any(example_value) and (
  2977. not tx.export or not config.capture_sparse_compute
  2978. ):
  2979. unimplemented(
  2980. gb_type="Attempted to wrap sparse Tensor with VariableTracker",
  2981. context=str(example_value),
  2982. explanation="torch.compile does not support sparse Tensors with VariableTracker",
  2983. hints=[*graph_break_hints.SPARSE_TENSOR],
  2984. )
  2985. var = construct_tensor_variable(
  2986. target_cls, tx, proxy, example_value, subclass_type, options
  2987. )
  2988. # NOTE: [Side effect tracking for newly constructed tensor]
  2989. # For newly constructed objects that have mutable attributes, we usually
  2990. # construct their VariableTracker via `track_object_new`, but since
  2991. # tensor variable construction is a bit different, we handle them
  2992. # specially here. This ensures that codegen will actually generate the
  2993. # attribute mutations on this tensor.
  2994. #
  2995. # NOTE we pass a dummy object as the `item` argument to avoid
  2996. # constructing a dummy _tensor_ object. The object isn't used for
  2997. # newly constructed VTs anyways.
  2998. assert isinstance(var, VariableTracker)
  2999. tx.output.side_effects._track_obj(
  3000. proxy, var, mutation_type_cls=AttributeMutationNew
  3001. )
  3002. return var
  3003. elif (
  3004. hasattr(proxy.node.target, "__name__")
  3005. and proxy.node.target.__name__ == "set_state"
  3006. # type: ignore[attr-defined]
  3007. and isinstance(proxy.node.target.__self__, torch._C.Generator)
  3008. or proxy.node.target is torch.random.set_rng_state
  3009. ):
  3010. assert type(proxy.node.target) is not str
  3011. # pyrefly: ignore[bad-argument-type]
  3012. return TorchInGraphFunctionVariable(proxy.node.target)
  3013. elif (
  3014. proxy.node.target is torch._C._DisableFuncTorch
  3015. or proxy.node.target is torch.cuda._is_in_bad_fork
  3016. ):
  3017. return UserDefinedObjectVariable(example_value)
  3018. elif istype(example_value, torch.Size) and all(
  3019. isinstance(x, int) for x in example_value
  3020. ):
  3021. sizes = [ConstantVariable.create(x) for x in example_value]
  3022. return SizeVariable(sizes, **options)
  3023. elif isinstance(example_value, (tuple, list)):
  3024. set_example_value(proxy.node, example_value)
  3025. unpacked = []
  3026. for i, val in enumerate(example_value):
  3027. if val is None:
  3028. # nn.MultiheadAttention() can return None, see issue #175
  3029. unpacked.append(
  3030. ConstantVariable.create(None, **options),
  3031. )
  3032. else:
  3033. proxy_i = proxy.tracer.create_proxy(
  3034. kind="call_function",
  3035. target=operator.getitem,
  3036. args=(proxy, i),
  3037. kwargs={},
  3038. )
  3039. if "source" in options:
  3040. # This path should only trigger for list stealing, so it's
  3041. # safe to use `GetItemSource`.
  3042. assert isinstance(example_value, list)
  3043. source = options["source"]
  3044. options_i = options.copy()
  3045. options_i["source"] = GetItemSource(
  3046. base=source,
  3047. index=i,
  3048. index_is_slice=False,
  3049. )
  3050. else:
  3051. # use the same options object as parent
  3052. options_i = options
  3053. # WARNING: this assumes the same target_cls as this tuple/list call
  3054. unpacked.append(
  3055. wrap_fx_proxy_cls(
  3056. # pyrefly: ignore[bad-argument-type]
  3057. target_cls=target_cls,
  3058. tx=tx,
  3059. proxy=proxy_i,
  3060. example_value=val,
  3061. **options_i,
  3062. )
  3063. )
  3064. if isinstance(example_value, torch.Size):
  3065. # NB: Keep the old proxy around. See SizeVariable for an
  3066. # explanation why
  3067. return SizeVariable(unpacked, proxy, **options)
  3068. elif istype(example_value, tuple):
  3069. return TupleVariable(unpacked, **options)
  3070. elif istype(example_value, (list, immutable_list)):
  3071. return ListVariable(unpacked, **options)
  3072. else:
  3073. assert (
  3074. example_value.__class__.__module__ == "torch.return_types"
  3075. or hasattr(example_value, "_fields")
  3076. ), (
  3077. f"expected {example_value.__class__.__module__} == torch.return_types or named tuple but got {type(example_value)}"
  3078. )
  3079. return NamedTupleVariable(unpacked, example_value.__class__, **options) # type: ignore[arg-type]
  3080. elif example_value is None or proxy.node.target is torch.manual_seed:
  3081. return ConstantVariable.create(None, **options)
  3082. elif isinstance(example_value, (torch.SymInt, torch.SymFloat, torch.SymBool)):
  3083. tx.output.current_tracer.track_produced_symints(example_value, proxy)
  3084. set_example_value(proxy.node, example_value)
  3085. return SymNodeVariable.create(tx, proxy, example_value, **options)
  3086. elif (
  3087. isinstance(example_value, torch.Stream)
  3088. and proxy.node.target is get_external_object_by_index
  3089. ) or proxy.node.target in [
  3090. device_interface.current_stream
  3091. for _, device_interface in get_registered_device_interfaces()
  3092. ]:
  3093. set_example_value(proxy.node, example_value)
  3094. index = None
  3095. if proxy.node.target is get_external_object_by_index:
  3096. index = proxy.node.args[0]
  3097. # type: ignore[arg-type]
  3098. return StreamVariable(proxy, example_value, index, **options)
  3099. elif (
  3100. isinstance(example_value, torch.Event)
  3101. and proxy.node.target is get_external_object_by_index
  3102. ) or proxy.node.target in [
  3103. device_interface.current_stream
  3104. for _, device_interface in get_registered_device_interfaces()
  3105. ]:
  3106. index = None
  3107. if proxy.node.target is get_external_object_by_index:
  3108. index = proxy.node.args[0]
  3109. set_example_value(proxy.node, example_value)
  3110. # type: ignore[arg-type]
  3111. return EventVariable(proxy, example_value, index, **options)
  3112. elif (
  3113. inspect.isclass(proxy.node.target)
  3114. and issubclass(proxy.node.target, torch.Event)
  3115. ) or proxy.node.target in [
  3116. device_interface.Event
  3117. for _, device_interface in get_registered_device_interfaces()
  3118. ]:
  3119. set_example_value(proxy.node, example_value)
  3120. return EventVariable(proxy, example_value, None, **options)
  3121. elif proxy.node.target == "query" and proxy.node.op == "call_method":
  3122. set_example_value(proxy.node, example_value)
  3123. return ConstantVariable(example_value, **options)
  3124. elif (
  3125. example_value is not None
  3126. and isinstance(example_value, torch.Event)
  3127. and proxy.node.target == "record_event"
  3128. and proxy.node.op == "call_method"
  3129. ):
  3130. set_example_value(proxy.node, example_value)
  3131. return EventVariable(proxy, example_value, None, **options)
  3132. elif isinstance(example_value, int) and (
  3133. proxy.node.target
  3134. in [
  3135. torch.sym_int,
  3136. getattr,
  3137. operator.getitem,
  3138. torch._utils._element_size,
  3139. torch.seed,
  3140. operator.mod,
  3141. torch._functorch.vmap._validate_and_get_batch_size,
  3142. torch._functorch.predispatch._vmap_increment_nesting,
  3143. torch._functorch.predispatch._vmap_decrement_nesting,
  3144. # some mac builds are missing torch.distributed.get_rank()
  3145. getattr(torch.distributed, "get_rank", _missing),
  3146. getattr(torch.distributed, "get_world_size", _missing),
  3147. # This always wants to be in the graph, even if the constraint
  3148. # results in a constant int
  3149. torch._constrain_as_size,
  3150. ]
  3151. or (
  3152. # TODO: this is a little sus, because we didn't check what the self is
  3153. proxy.node.op == "call_method" and proxy.node.target == "bit_length"
  3154. )
  3155. ):
  3156. set_example_value(proxy.node, example_value)
  3157. return ConstantVariable.create(example_value, **options)
  3158. elif isinstance(example_value, torch.backends.cuda.SDPAParams):
  3159. from .sdpa import SDPAParamsVariable
  3160. set_example_value(proxy.node, example_value)
  3161. return SDPAParamsVariable(proxy, **options)
  3162. elif isinstance(example_value, bool) and (
  3163. proxy.node.target
  3164. in [
  3165. torch._C._are_functorch_transforms_active,
  3166. torch._C._functorch.is_batchedtensor,
  3167. torch.backends.cuda.is_flash_attention_available,
  3168. torch.backends.cuda.can_use_flash_attention,
  3169. torch.backends.cuda.can_use_efficient_attention,
  3170. torch._C._get_cudnn_sdp_enabled,
  3171. torch._C._get_flash_sdp_enabled,
  3172. torch._C._get_mem_efficient_sdp_enabled,
  3173. torch._C._get_math_sdp_enabled,
  3174. torch._C._get_overrideable_sdp_enabled,
  3175. "is_integer",
  3176. ]
  3177. + list(supported_const_comparison_op_values.keys())
  3178. ):
  3179. set_example_value(proxy.node, example_value)
  3180. return ConstantVariable.create(example_value, **options)
  3181. elif isinstance(example_value, (int, float, bool)) and (
  3182. proxy.node.target is call_torchbind
  3183. or proxy.node.target is flat_apply
  3184. or (proxy.node.op == "call_method" and proxy.node.target == "item")
  3185. ):
  3186. set_example_value(proxy.node, example_value)
  3187. return ConstantVariable.create(example_value, **options)
  3188. elif isinstance(example_value, float) or proxy.node.target in ["hex", "__round__"]:
  3189. set_example_value(proxy.node, example_value)
  3190. return ConstantVariable.create(example_value, **options)
  3191. elif is_opaque_type(type(example_value)):
  3192. # This is for handling opaque objects in custom ops
  3193. if is_opaque_value_type(type(example_value)):
  3194. proxy = example_value # pyrefly: ignore[bad-assignment]
  3195. fake_script_obj = torch._library.fake_class_registry.maybe_to_fake_obj(
  3196. tx.output.fake_mode, example_value
  3197. )
  3198. return TorchScriptObjectVariable.create(
  3199. proxy,
  3200. fake_script_obj,
  3201. )
  3202. else:
  3203. unimplemented(
  3204. gb_type="torch.* op returned non-Tensor",
  3205. context=f"example_value type: {typestr(example_value)}; op: {proxy.node.op}; target: {proxy.node.target}",
  3206. explanation="torch.* ops that return a non-Tensor cannot be traced into the Dynamo FX graph output",
  3207. hints=[],
  3208. )
  3209. def infer_subclass_type(value: T) -> type[T] | None:
  3210. if type(value) in (
  3211. torch.Tensor,
  3212. torch.nn.Parameter,
  3213. torch._subclasses.fake_tensor.FakeTensor,
  3214. torch._subclasses.functional_tensor.FunctionalTensor,
  3215. ) or is_traceable_wrapper_subclass(value):
  3216. # Ordinarily, we would fakeify a tensor so that it can get dynamic
  3217. # shapes and be computed on without triggering actual operations.
  3218. # However, how can we fakeify a tensor subclass? Ordinary
  3219. # inheritance (nor multiple inheritance) won't work work.
  3220. #
  3221. # Instead, our plan is to *manually simulate* the tensor subclass
  3222. # inheriting from a fake tensor with dynamo. This means our
  3223. # data representation for a tensor subclass will be a fake tensor
  3224. # + tensor subclass type + any extra data the subclass may have
  3225. # been storing on the tensor. Because all Python accesses are
  3226. # mediated through TensorWithTFOverrideVariable, we can ensure
  3227. # that we dispatch differently, e.g., according to
  3228. # __torch_function__
  3229. #
  3230. # To simplify things for now, the __dict__ tracking bits haven't
  3231. # been implemented yet, but they can be added into this design at
  3232. # a later point in time.
  3233. return None
  3234. else:
  3235. return type(value)
  3236. def get_specialized_props(
  3237. target_cls: Any,
  3238. tx: "InstructionTranslatorBase",
  3239. example_value: Any,
  3240. subclass_type: type | None,
  3241. ) -> dict[str, Any]:
  3242. specialized_props = target_cls.specialize(example_value)
  3243. # TODO: not sure about this fake mode test
  3244. if (
  3245. isinstance(example_value, torch._subclasses.fake_tensor.FakeTensor)
  3246. and example_value.fake_mode is tx.fake_mode
  3247. ):
  3248. if subclass_type:
  3249. tensor_type = subclass_type
  3250. elif isinstance(example_value, torch.nn.Parameter):
  3251. tensor_type = torch.nn.Parameter
  3252. elif isinstance(example_value, torch.nn.Buffer):
  3253. tensor_type = torch.nn.Buffer
  3254. else:
  3255. tensor_type = torch.Tensor
  3256. specialized_props["class_type"] = tensor_type
  3257. return specialized_props
  3258. def construct_tensor_variable(
  3259. target_cls: type[VTTypeAlias],
  3260. tx: "InstructionTranslatorBase",
  3261. proxy: torch.fx.Proxy,
  3262. example_value: Any,
  3263. subclass_type: type | None,
  3264. options: dict[str, Any],
  3265. ) -> VTTypeAlias:
  3266. """
  3267. Actually construct a tensor variable after all the pre-processing from
  3268. wrapping a pre-existing or newly created tensor value.
  3269. """
  3270. # NB: In most (all?) cases, this does not actually do a clone.
  3271. # (WARNING: this means that if we mutate metadata on the fake
  3272. # tensor, the stored example value will update too!)
  3273. example_value = _clone_input(example_value, tx.fake_mode)
  3274. set_example_value(proxy.node, example_value)
  3275. # We bind the unbacked symints in sizes/trdies of tensor lazily.
  3276. # So that subgraphs can access the unbacked symbol's proxy in parent graph
  3277. # when lifting unbacked symbols of input tensors to subgraph inputs.
  3278. # We do it lazily because the tensor may not be used in subgraphs.
  3279. if proxy.node.op != "placeholder":
  3280. tx.output.current_tracer.track_produced_symints(example_value, proxy)
  3281. options.update(get_specialized_props(target_cls, tx, example_value, subclass_type))
  3282. return target_cls(proxy, **options)
  3283. def get_automatic_dynamic_shapes_mark_as() -> DimDynamic:
  3284. if config.automatic_dynamic_shapes_mark_as == "dynamic":
  3285. return DimDynamic.DYNAMIC
  3286. elif config.automatic_dynamic_shapes_mark_as == "unbacked":
  3287. return DimDynamic.UNBACKED
  3288. else:
  3289. raise ValueError(
  3290. f"invalid automatic_dynamic_shapes_mark_as = {config.automatic_dynamic_shapes_mark_as}"
  3291. )
  3292. _DYNAMIC_SOURCES: set[str] | None = None
  3293. _DYNAMIC_SOURCES_CONFIG_HASH: int | None = None
  3294. def get_dynamic_sources() -> set[str]:
  3295. global _DYNAMIC_SOURCES, _DYNAMIC_SOURCES_CONFIG_HASH
  3296. current_hash = hash(torch.compiler.config.dynamic_sources)
  3297. # If we have already calculated the sources and the config hasn't changed, return cached result
  3298. if _DYNAMIC_SOURCES is not None and _DYNAMIC_SOURCES_CONFIG_HASH == current_hash:
  3299. return _DYNAMIC_SOURCES
  3300. # Config has changed or first time, (re)calculate the sources
  3301. _DYNAMIC_SOURCES = {
  3302. s
  3303. for s in torch.compiler.config.dynamic_sources.replace(" ", "").split(",")
  3304. if s
  3305. }
  3306. _DYNAMIC_SOURCES_CONFIG_HASH = current_hash
  3307. return _DYNAMIC_SOURCES
  3308. def is_dynamic_source(source_name: str) -> bool:
  3309. dynamic_sources = get_dynamic_sources()
  3310. for pattern in dynamic_sources:
  3311. if pattern == source_name or re.match(pattern, source_name):
  3312. log.debug(
  3313. "%s was marked dynamic due to dynamic source allowlist pattern: %s",
  3314. source_name,
  3315. pattern,
  3316. )
  3317. return True
  3318. return False
  3319. def record_automatic_dynamic(
  3320. tx: "InstructionTranslatorBase", name: str, e: torch.Tensor
  3321. ) -> FrameStateSizeEntry:
  3322. # This mimics stride inference algorithm in _create_symbolic_sizes_strides_storage_offset
  3323. ex_size = e.size()
  3324. if not is_sparse_any(e):
  3325. ex_stride = e.stride()
  3326. dim = e.dim()
  3327. stride = [None] * dim
  3328. pending = [(ex_stride[i], -i) for i in range(dim)]
  3329. pending.sort(key=_nested_int_aware_sort)
  3330. candidates = {}
  3331. for i_stride, neg_i in pending:
  3332. i = -neg_i
  3333. # pyrefly: ignore [unsupported-operation]
  3334. stride[i] = candidates.get(i_stride, i_stride)
  3335. # pyrefly: ignore [no-matching-overload]
  3336. candidates.setdefault(i_stride * ex_size[i], InferStride(i))
  3337. else:
  3338. # pyrefly: ignore [implicit-any]
  3339. stride = []
  3340. return process_automatic_dynamic(
  3341. # type: ignore[arg-type]ks
  3342. tx,
  3343. name,
  3344. # type: ignore[arg-type]
  3345. FrameStateSizeEntry.make_tensor(tuple(ex_size), tuple(stride)),
  3346. )
  3347. _UNBACKED_SOURCES: set[str] | None = None
  3348. _UNBACKED_SOURCES_CONFIG_HASH: int | None = None
  3349. def get_unbacked_sources() -> set[str]:
  3350. global _UNBACKED_SOURCES, _UNBACKED_SOURCES_CONFIG_HASH
  3351. current_hash = hash(torch.compiler.config.unbacked_sources)
  3352. # If we have already calculated the sources and the config hasn't changed, return cached result
  3353. if _UNBACKED_SOURCES is not None and _UNBACKED_SOURCES_CONFIG_HASH == current_hash:
  3354. return _UNBACKED_SOURCES
  3355. # Config has changed or first time, (re)calculate the sources
  3356. _UNBACKED_SOURCES = {
  3357. s
  3358. for s in torch.compiler.config.unbacked_sources.replace(" ", "").split(",")
  3359. if s
  3360. }
  3361. _UNBACKED_SOURCES_CONFIG_HASH = current_hash
  3362. return _UNBACKED_SOURCES
  3363. def is_unbacked_source(source_name: str) -> bool:
  3364. unbacked_sources = get_unbacked_sources()
  3365. for pattern in unbacked_sources:
  3366. if pattern == source_name or re.match(pattern, source_name):
  3367. log.debug(
  3368. "%s was marked unbacked due to unbacked source allowlist pattern: %s",
  3369. source_name,
  3370. pattern,
  3371. )
  3372. return True
  3373. return False
  3374. # Performs automatic dynamic dim determination.
  3375. # Returns a SymbolicContext
  3376. def _automatic_dynamic(
  3377. e: Any,
  3378. tx: "InstructionTranslatorBase",
  3379. source: Source,
  3380. static_shapes: bool,
  3381. outer_only: bool = False,
  3382. ) -> SymbolicContext:
  3383. # strided NT not supported
  3384. if e.is_nested and not isinstance(
  3385. e, torch.nested._internal.nested_tensor.NestedTensor
  3386. ):
  3387. unimplemented(
  3388. gb_type="Encountered strided NestedTensor in automatic dynamic dim determination",
  3389. context="",
  3390. explanation="torch.compile does not support strided NestedTensor",
  3391. hints=[],
  3392. )
  3393. name = source.name
  3394. prior_policy = tx.output.tracing_context.tensor_to_context.get(e, None)
  3395. shape_env_to_source_to_symbol_cache = (
  3396. prior_policy.shape_env_to_source_to_symbol_cache if prior_policy else {}
  3397. )
  3398. # Get base context if the tensor is a view
  3399. view_base_context: SymbolicContext | None = None
  3400. if e._is_view():
  3401. base_source = AttrSource(source, "_base")
  3402. view_base_context = _automatic_dynamic(e._base, tx, base_source, static_shapes)
  3403. if is_traceable_wrapper_subclass(e) and not outer_only:
  3404. # Get symbolic context for outer tensor
  3405. outer_context = _automatic_dynamic(
  3406. e, tx, source, static_shapes, outer_only=True
  3407. )
  3408. assert isinstance(outer_context, StatefulSymbolicContext)
  3409. # Get symbolic contexts for inner tensors
  3410. inner_contexts = {} # mapping from attr -> symbolic context
  3411. attrs, _ = type(e).__tensor_flatten__(e)
  3412. for attr in attrs:
  3413. inner_tensor = getattr(e, attr)
  3414. inner_source = AttrSource(source, attr)
  3415. inner_contexts[attr] = _automatic_dynamic(
  3416. inner_tensor, tx, inner_source, static_shapes
  3417. )
  3418. return SubclassSymbolicContext(
  3419. dynamic_sizes=outer_context.dynamic_sizes,
  3420. dynamic_strides=outer_context.dynamic_strides,
  3421. constraint_sizes=outer_context.constraint_sizes,
  3422. constraint_strides=outer_context.constraint_strides,
  3423. view_base_context=view_base_context,
  3424. tensor_source=outer_context.tensor_source,
  3425. shape_env_to_source_to_symbol_cache=outer_context.shape_env_to_source_to_symbol_cache,
  3426. inner_contexts=inner_contexts,
  3427. )
  3428. if static_shapes and not is_dynamic_source(name):
  3429. return StatefulSymbolicContext(
  3430. dynamic_sizes=[DimDynamic.STATIC] * e.dim(),
  3431. dynamic_strides=[DimDynamic.INFER_STRIDE] * e.dim(),
  3432. constraint_sizes=[None] * e.dim(),
  3433. constraint_strides=[None] * e.dim(),
  3434. view_base_context=view_base_context,
  3435. tensor_source=source,
  3436. shape_env_to_source_to_symbol_cache=shape_env_to_source_to_symbol_cache,
  3437. )
  3438. # We preserve the dynamism of inputs. For example, when users call
  3439. # make_fx(torch.cond, tracing_mode="symbolic")(*args), inputs have SymInt sizes.
  3440. from torch.fx.experimental.symbolic_shapes import is_nested_int
  3441. if any(isinstance(s, SymInt) and not is_nested_int(s) for s in e.size()):
  3442. return StatefulSymbolicContext(
  3443. dynamic_sizes=[
  3444. DimDynamic.DYNAMIC if isinstance(s, SymInt) else DimDynamic.STATIC
  3445. for s in e.size()
  3446. ],
  3447. dynamic_strides=[DimDynamic.INFER_STRIDE] * e.dim(),
  3448. constraint_sizes=[None] * e.dim(),
  3449. constraint_strides=[None] * e.dim(),
  3450. view_base_context=view_base_context,
  3451. tensor_source=source,
  3452. shape_env_to_source_to_symbol_cache=shape_env_to_source_to_symbol_cache,
  3453. )
  3454. # Prep for automatic dynamic
  3455. frame_state_entry = record_automatic_dynamic(tx, name, e)
  3456. # TODO: index export_constraints ahead of time so we don't have to
  3457. # do a linear scan every time here
  3458. t_id = id(e)
  3459. # pyrefly: ignore [implicit-any]
  3460. dim2constraint = {}
  3461. def update_dim2constraint(
  3462. dim: int, constraint_range: "StrictMinMaxConstraint", name: str
  3463. ) -> None:
  3464. if dim in dim2constraint:
  3465. from torch.fx.experimental.symbolic_shapes import StrictMinMaxConstraint
  3466. old_constraint_range, old_name = dim2constraint[dim]
  3467. new_constraint_range = StrictMinMaxConstraint(
  3468. vr=constraint_range.vr & old_constraint_range.vr,
  3469. warn_only=False,
  3470. )
  3471. # It is possible for (non-None) old_name and name to be different
  3472. # but this will only happen the corresponding Dims can be derived equal.
  3473. new_name = old_name or name
  3474. dim2constraint[dim] = new_constraint_range, new_name
  3475. else:
  3476. dim2constraint[dim] = constraint_range, name
  3477. from torch.export.dynamic_shapes import _RelaxedConstraint
  3478. if tx.output.export_constraints is not None:
  3479. # type: ignore[iterable]
  3480. for constraint in tx.output.export_constraints:
  3481. if isinstance(constraint, _RelaxedConstraint):
  3482. continue
  3483. if constraint.t_id == t_id:
  3484. update_dim2constraint(
  3485. constraint.dim, constraint.constraint_range, constraint.name
  3486. )
  3487. dynamic_sizes = []
  3488. dynamic_strides = []
  3489. constraint_sizes = []
  3490. constraint_strides = []
  3491. specialize_on = []
  3492. for i in range(e.dim()):
  3493. # NB: mark dynamic has precedence over static
  3494. marked_strict_unbacked = i in getattr(
  3495. e, "_dynamo_strict_unbacked_indices", set()
  3496. )
  3497. marked_unbacked = i in getattr(e, "_dynamo_unbacked_indices", set())
  3498. marked_dynamic = i in getattr(e, "_dynamo_dynamic_indices", set())
  3499. marked_weak_dynamic = i in getattr(e, "_dynamo_weak_dynamic_indices", set())
  3500. marked_static = i in getattr(e, "_dynamo_static_indices", set())
  3501. specialize_on.append(getattr(e, "_specialize_on", {}).get(i, []))
  3502. # Reflect the user directive in the frame_state
  3503. # For dynamic, apply None always
  3504. normalized_source_name = normalize_source_name(source.name)
  3505. base_source = source
  3506. if isinstance(base_source, ChainedSource):
  3507. base_source = base_source.get_base()
  3508. if marked_dynamic or (
  3509. isinstance(base_source, LocalSource)
  3510. and base_source.dynamism is not None
  3511. # pyrefly: ignore[no-matching-overload]
  3512. and dict(base_source.dynamism).get(normalized_source_name, {i: False})[i]
  3513. ):
  3514. # TODO: This can be batched
  3515. # TODO: Doing this here is kind of sus, maybe better to set this
  3516. # up when we initially created the FrameStateSizeEntry to bong
  3517. # into the mutable state
  3518. log.debug("automatic dynamic %s marked dynamic", name)
  3519. mark_size = [auto_unset] * e.dim()
  3520. # pyrefly: ignore [unsupported-operation]
  3521. mark_size[i] = auto_dynamic
  3522. # pyrefly: ignore [bad-argument-type]
  3523. frame_state_entry |= FrameStateSizeEntry.make_size(size=mark_size)
  3524. # NB: both static and dynamic have precedence over
  3525. automatic_dynamic_size = (
  3526. config.automatic_dynamic_shapes and frame_state_entry.is_size_dynamic(i)
  3527. )
  3528. # NB: previously, if size was dynamic, we wouldn't make its stride
  3529. # dynamic. But now, because of InferStride concept, we will properly
  3530. # not make stride dynamic even if it's wobbling
  3531. automatic_dynamic_stride = (
  3532. config.automatic_dynamic_shapes and frame_state_entry.is_stride_dynamic(i)
  3533. )
  3534. if is_dynamic_source(name):
  3535. log.debug("%s marked dynamic via source whitelist", name)
  3536. automatic_dynamic_size = True
  3537. if is_unbacked_source(name):
  3538. log.debug("%s marked unbacked via source whitelist", name)
  3539. automatic_dynamic_size = True
  3540. automatic_dynamic = automatic_dynamic_size or automatic_dynamic_stride
  3541. # We will process constraints first, as they will imply that we
  3542. # have a dynamic dimension
  3543. # Precedence: export constraints > eager constraints
  3544. constraint = dim2constraint.get(i)
  3545. if constraint is None:
  3546. constraint_size = None
  3547. constraint_stride = None
  3548. if marked_dynamic and not config.allow_ignore_mark_dynamic:
  3549. # constraint_stride is deliberaly kept None because no easy way to provide value ranges for mark dynamic
  3550. constraint_stride = None
  3551. if hasattr(e, "_dynamo_dynamic_range"):
  3552. dim_range = [
  3553. dr for dr in e._dynamo_dynamic_range if dr.dim == i
  3554. ].pop()
  3555. if dim_range.min is None and dim_range.max is None:
  3556. constraint_size = RelaxedUnspecConstraint(warn_only=False)
  3557. else:
  3558. from torch.fx.experimental.symbolic_shapes import (
  3559. StrictMinMaxConstraint,
  3560. )
  3561. constraint_size = StrictMinMaxConstraint(
  3562. vr=ValueRanges(lower=dim_range.min, upper=dim_range.max),
  3563. warn_only=False,
  3564. )
  3565. else:
  3566. constraint_size = RelaxedUnspecConstraint(warn_only=False)
  3567. elif marked_strict_unbacked:
  3568. constraint_size = RelaxedUnspecConstraint(warn_only=False)
  3569. elif not marked_static and automatic_dynamic:
  3570. set_feature_use("dynamo.automatic_dynamic_shapes", True)
  3571. if automatic_dynamic_size:
  3572. constraint_size = RelaxedUnspecConstraint(warn_only=True)
  3573. if automatic_dynamic_stride:
  3574. constraint_stride = RelaxedUnspecConstraint(warn_only=True)
  3575. else:
  3576. if not marked_static and not config.automatic_dynamic_shapes:
  3577. set_feature_use("dynamo.automatic_dynamic_shapes", False)
  3578. constraint_size = None
  3579. constraint_stride = None
  3580. else:
  3581. constraint_size, name_ = constraint
  3582. constraint_stride = None
  3583. dim_name = f"{name}.size()[{i}]"
  3584. tx.output.shape_env.source_name_to_debug_name[dim_name] = name_
  3585. constraint_sizes.append(constraint_size)
  3586. constraint_strides.append(constraint_stride)
  3587. if marked_unbacked or is_unbacked_source(name):
  3588. dynamic_size = DimDynamic.UNBACKED
  3589. elif (
  3590. constraint_size is not None
  3591. or marked_dynamic
  3592. or marked_weak_dynamic
  3593. or is_nested_int(e.size()[i])
  3594. ):
  3595. # NB: We could assert static_shapes is False here, but it
  3596. # seems better to allow the user to override symbolic_context in this
  3597. # case
  3598. if automatic_dynamic:
  3599. dynamic_size = get_automatic_dynamic_shapes_mark_as()
  3600. else:
  3601. dynamic_size = DimDynamic.DYNAMIC
  3602. elif static_shapes or config.assume_static_by_default or marked_static:
  3603. dynamic_size = DimDynamic.STATIC
  3604. else:
  3605. # TODO: When does this show up?
  3606. dynamic_size = DimDynamic.DUCK
  3607. if constraint_stride is not None:
  3608. dynamic_stride = DimDynamic.DYNAMIC
  3609. else:
  3610. dynamic_stride = DimDynamic.INFER_STRIDE
  3611. dynamic_sizes.append(dynamic_size)
  3612. dynamic_strides.append(dynamic_stride)
  3613. return StatefulSymbolicContext(
  3614. dynamic_sizes=dynamic_sizes,
  3615. dynamic_strides=dynamic_strides,
  3616. constraint_sizes=constraint_sizes,
  3617. # pyrefly: ignore [bad-argument-type]
  3618. constraint_strides=constraint_strides,
  3619. specialize_on=specialize_on,
  3620. view_base_context=view_base_context,
  3621. tensor_source=source,
  3622. shape_env_to_source_to_symbol_cache=shape_env_to_source_to_symbol_cache,
  3623. shape_ids=getattr(e, "_dynamo_shape_ids", None),
  3624. )
  3625. # See note [Tensor Fakification and Symbol Caching]
  3626. def wrap_to_fake_tensor_and_record(
  3627. e: Any,
  3628. tx: "InstructionTranslatorBase",
  3629. *,
  3630. source: Source | None,
  3631. is_tensor: bool,
  3632. parent_context: Any | None = None,
  3633. ) -> Any:
  3634. _t0 = time.time_ns()
  3635. try:
  3636. return _wrap_to_fake_tensor_and_record_impl(
  3637. e, tx, source=source, is_tensor=is_tensor, parent_context=parent_context
  3638. )
  3639. finally:
  3640. tx.output.bytecode_tracing_timings.wrap_to_fake_tensor_and_record_ns += (
  3641. time.time_ns() - _t0
  3642. )
  3643. def _wrap_to_fake_tensor_and_record_impl(
  3644. e: Any,
  3645. tx: "InstructionTranslatorBase",
  3646. *,
  3647. source: Source | None,
  3648. is_tensor: bool,
  3649. parent_context: Any | None = None,
  3650. ) -> Any:
  3651. if (
  3652. type(e) in (torch.Tensor, torch.nn.Parameter, FakeTensor)
  3653. or isinstance(e, torch.Tensor)
  3654. or is_traceable_wrapper_subclass(e)
  3655. ):
  3656. assert source is not None
  3657. static_shapes, _reason = tensor_always_has_static_shape(
  3658. e,
  3659. is_tensor,
  3660. tensor_source=source,
  3661. )
  3662. if not parent_context:
  3663. symbolic_context = _automatic_dynamic(e, tx, source, static_shapes)
  3664. else:
  3665. # Parent contexts are passed in when we are recursively creating
  3666. # fake tensors for subclasses. A better design would be not to create a
  3667. # parent/child relationship, but to recursively call _automatic_dynamic
  3668. # as we recursively call wrap_to_fake_tensor_and_record. This runs
  3669. # into bugs around how meta_utils knows and works to create fake tensors
  3670. # with tensor subclasses. Ideally, dynamo would drive both the recursive
  3671. # wrap_to_fake_tensor_and_record and _automatic_dynamic policy creation.
  3672. assert isinstance(source, AttrSource)
  3673. inner_context_name = source.member
  3674. symbolic_context = parent_context.inner_contexts[inner_context_name]
  3675. log.debug(
  3676. "wrap_to_fake %s %s %s %s",
  3677. source.name,
  3678. tuple(e.shape),
  3679. symbolic_context,
  3680. type(e),
  3681. )
  3682. # Note [enable_python_dispatcher in dynamo]
  3683. # Dynamo disables itself when it runs fake tensor prop, which means that tensor subclasses
  3684. # have no way to know (purely based off of global state) if they are currently being run under compile or not.
  3685. # we use enable_python_dispatcher mainly to tweak the DispatchKeyState so that subclass authors
  3686. # can check it to know if they are running in an eager context or not
  3687. with enable_python_dispatcher():
  3688. assert tx.fake_mode is not None
  3689. fake_e = wrap_fake_exception(
  3690. lambda: tx.fake_mode.from_tensor(
  3691. e, # type: ignore[arg-type]
  3692. source=source,
  3693. symbolic_context=symbolic_context,
  3694. )
  3695. )
  3696. if (
  3697. source is not None
  3698. and isinstance(fake_e, FakeTensor)
  3699. and (sym_val := fake_e.item_memo) is not None
  3700. ):
  3701. tx.output.tracked_fakes.append(
  3702. TrackedFake(sym_val, CallMethodItemSource(source), symbolic_context)
  3703. )
  3704. if is_traceable_wrapper_subclass(fake_e):
  3705. attrs, _ = fake_e.__tensor_flatten__()
  3706. for attr in attrs:
  3707. fake_inner = getattr(fake_e, attr)
  3708. inner = getattr(e, attr)
  3709. inner_source = AttrSource(source, attr)
  3710. wrap_to_fake_tensor_and_record(
  3711. inner,
  3712. tx,
  3713. source=inner_source,
  3714. is_tensor=isinstance(fake_inner, torch.Tensor),
  3715. parent_context=symbolic_context,
  3716. )
  3717. tx.output.tracing_context.tensor_to_context[e] = symbolic_context
  3718. if is_sparse_any(fake_e):
  3719. # TODO: for TensorGuards, this eventually may need more
  3720. # fields for the size/stride of any other constituents
  3721. values = fake_e._values() if fake_e.is_sparse else fake_e.values()
  3722. tx.output.input_source_to_sizes_strides[source] = {
  3723. "size": fake_e.size(),
  3724. # TODO: revise this, but for now this stride instead of ()
  3725. # avoids SegFault with PYTORCH_TEST_WITH_DYNAMO=1
  3726. "stride": (1,) * fake_e.ndim,
  3727. "values_size": values.size(),
  3728. "values_stride": values.stride(),
  3729. }
  3730. else:
  3731. tx.output.input_source_to_sizes_strides[source] = {
  3732. "size": fake_e.size(),
  3733. "stride": fake_e.stride(),
  3734. }
  3735. if (
  3736. is_tensor
  3737. and not (static_shapes and source.is_specialized_nn_module())
  3738. and not is_constant_source(source)
  3739. ):
  3740. tx.output.tracked_fakes.append(
  3741. TrackedFake(fake_e, source, symbolic_context)
  3742. )
  3743. tx.output.tracked_fakes_id_to_source[id(e)].append(source)
  3744. return fake_e
  3745. else:
  3746. return e
  3747. class SourcelessBuilder:
  3748. """
  3749. Like builder, but stateless and does not require a source. Useful for simple type->VT objects, or objects
  3750. that are being created/evaporated during inlining (ex: consider a locally made list of tensors we then iterate over
  3751. .), such a list should not show up as an artifact from inputs, nor in reconstruction, nor in the graph. However,
  3752. there may be reasons to represent it as a ListVariable internally.
  3753. NOTE - Objects produced here are born UNGUARDED due to the nature of sources!
  3754. NOTE - This class is very new! It will have some rough edges, but it was created to stem the bleeding of giant
  3755. if/else type->VariableTracker trees that were cropping up all over dynamo.
  3756. """
  3757. def __init__(self) -> None:
  3758. raise AssertionError("Use SourcelessBuilder.create()")
  3759. @overload
  3760. @staticmethod
  3761. def create(
  3762. tx: "InstructionTranslatorBase",
  3763. value: type[set[Any]]
  3764. | type[dict[Any, Any]]
  3765. | type[tuple[Any, ...]]
  3766. | type[list[Any]],
  3767. ) -> BuiltinVariable: ...
  3768. @overload
  3769. @staticmethod
  3770. def create(tx: "InstructionTranslatorBase", value: list[Any]) -> ListVariable: ...
  3771. @overload
  3772. @staticmethod
  3773. def create(
  3774. tx: "InstructionTranslatorBase", value: tuple[Any, ...]
  3775. ) -> TupleVariable: ...
  3776. @overload
  3777. @staticmethod
  3778. def create(
  3779. tx: "InstructionTranslatorBase", value: bool | int | float | str
  3780. ) -> ConstantVariable: ...
  3781. @overload
  3782. @staticmethod
  3783. def create(tx: "InstructionTranslatorBase", value: Any) -> VariableTracker: ...
  3784. @staticmethod
  3785. def create(tx: "InstructionTranslatorBase", value: Any) -> VariableTracker:
  3786. value_type = type(value)
  3787. # type: ignore[attr-defined]
  3788. fast_handler = SourcelessBuilder._type_handlers.get(value_type)
  3789. if fast_handler:
  3790. return fast_handler(tx, value)
  3791. if isinstance(value, VariableTracker):
  3792. # This is always valid to call, and useful for recursive calls.
  3793. return value
  3794. elif is_opaque_value_type(type(value)):
  3795. # This is for handling opaque objects in custom ops
  3796. fake_script_obj = torch._library.fake_class_registry.maybe_to_fake_obj(
  3797. tx.output.fake_mode, value
  3798. )
  3799. return TorchScriptObjectVariable.create(
  3800. value,
  3801. fake_script_obj,
  3802. )
  3803. # type: ignore[attr-defined]
  3804. elif isinstance(value, dataclasses._HAS_DEFAULT_FACTORY_CLASS):
  3805. return UserDefinedObjectVariable(value)
  3806. elif ConstantVariable.is_literal(value):
  3807. return ConstantVariable.create(value)
  3808. elif callable(value) and trace_rules.lookup_callable(value) is not None:
  3809. if trace_rules.is_callable_allowed(value):
  3810. tx.output.has_user_defined_allowed_in_graph = True
  3811. # pyrefly: ignore[not-callable, bad-argument-count]
  3812. return trace_rules.lookup_callable(value)(value)
  3813. elif callable(value) and UserDefinedClassVariable.is_supported_new_method(
  3814. value
  3815. ):
  3816. # NamedTuple._make uses an alias of tuple.__new__
  3817. # pyrefly: ignore[not-callable, bad-argument-count, missing-attribute]
  3818. obj = trace_rules.lookup_callable(value.__self__)(value.__self__)
  3819. return GetAttrVariable(obj, "__new__")
  3820. elif is_function_or_wrapper(value):
  3821. # pyrefly: ignore[not-callable, bad-argument-count]
  3822. return trace_rules.lookup(value)(value)
  3823. elif isinstance(
  3824. value, (enum.Enum, torch.DispatchKey, torch._C._functorch.TransformType)
  3825. ):
  3826. return EnumVariable(value)
  3827. elif isinstance(value, (type, abc.ABCMeta)):
  3828. if isinstance(value, type) and issubclass(value, enum.Enum):
  3829. return UserDefinedEnumClassVariable(value)
  3830. return UserDefinedClassVariable(value)
  3831. elif isinstance(value, types.MethodWrapperType):
  3832. return MethodWrapperVariable(value)
  3833. elif (
  3834. isinstance(value, types.MethodType)
  3835. # We only want to support sourceless class objects here
  3836. # An instance variable is not allowed and it should have source
  3837. and isinstance(value.__self__, (type, abc.ABCMeta))
  3838. ):
  3839. # value is a classmethod
  3840. assert getattr(value.__self__, value.__func__.__name__) == value
  3841. cls_obj_vt = SourcelessBuilder.create(tx, value.__self__)
  3842. try:
  3843. # pyrefly: ignore[bad-argument-type]
  3844. return cls_obj_vt.var_getattr(tx, value.__func__.__name__)
  3845. except NotImplementedError:
  3846. pass # failthrough to unimplemented branch
  3847. elif isinstance(value, torch.fx.graph_module.GraphModule):
  3848. return SourcelessGraphModuleVariable(value)
  3849. elif isinstance(value, torch.utils._pytree.TreeSpec):
  3850. return UserDefinedObjectVariable(value)
  3851. elif PlacementVariable.is_placement(value):
  3852. return PlacementVariable(value)
  3853. elif DeviceMeshVariable.is_device_mesh(value):
  3854. return DeviceMeshVariable(value)
  3855. elif value is functools.wraps:
  3856. return FunctoolsWrapsVariable(value)
  3857. elif isinstance(value, re.Pattern):
  3858. return ConstantLikeVariable(value)
  3859. elif isinstance(value, torch._dynamo.variables.lazy.LazySymNodeFormatString):
  3860. return ConstantVariable.create(str(value))
  3861. elif isinstance(value, type(torch._higher_order_ops.flex_attention_backward)):
  3862. return torch._dynamo.variables.higher_order_ops.FlexAttentionBackwardHighOrderVariable(
  3863. value
  3864. )
  3865. elif isinstance(value, (types.GenericAlias, types.UnionType)):
  3866. return TypingVariable(value)
  3867. elif is_namedtuple(value):
  3868. output = [
  3869. SourcelessBuilder.create(tx, getattr(value, name))
  3870. for name in namedtuple_fields(type(value))
  3871. ]
  3872. return NamedTupleVariable(output, tuple_cls=type(value))
  3873. elif (
  3874. isinstance(value, torch.SymInt)
  3875. and value.node.expr in tx.output.bound_symbols
  3876. ):
  3877. proxy = tx.output.bound_symbols[value.node.expr]
  3878. return SymNodeVariable.create(tx, proxy)
  3879. elif istype(value, object):
  3880. return ObjectVariable(value)
  3881. unimplemented(
  3882. gb_type="Unexpected type in sourceless builder",
  3883. context=f"{value_type.__module__}.{value_type.__qualname__}",
  3884. explanation=f"SourcelessBuilder.create does not know how to wrap {value_type}",
  3885. hints=[*graph_break_hints.DYNAMO_BUG],
  3886. )
  3887. @staticmethod
  3888. def wrap_constant_literal(value: object) -> VariableTracker:
  3889. assert ConstantVariable.is_literal(value)
  3890. return ConstantVariable.create(value=value)
  3891. @staticmethod
  3892. def make_type_handlers() -> dict[
  3893. type, Callable[["InstructionTranslator", Any], VariableTracker]
  3894. ]:
  3895. create = SourcelessBuilder.create
  3896. handlers: dict[
  3897. type, Callable[[InstructionTranslator, Any], VariableTracker]
  3898. ] = {}
  3899. for t in common_constant_types:
  3900. handlers[t] = lambda tx, value: ConstantVariable(value)
  3901. handlers[set] = lambda tx, value: SetVariable(
  3902. [create(tx, x) for x in value], mutation_type=ValueMutationNew()
  3903. )
  3904. handlers[OrderedSet] = lambda tx, value: OrderedSetVariable(
  3905. [create(tx, x) for x in value], mutation_type=ValueMutationNew()
  3906. )
  3907. handlers[dict] = lambda tx, value: ConstDictVariable(
  3908. {create(tx, k): create(tx, v) for k, v in value.items()},
  3909. type(value),
  3910. mutation_type=ValueMutationNew(),
  3911. )
  3912. handlers[list] = lambda tx, value: ListVariable(
  3913. [create(tx, x) for x in value], mutation_type=ValueMutationNew()
  3914. )
  3915. handlers[tuple] = lambda tx, value: TupleVariable(
  3916. [create(tx, x) for x in value]
  3917. )
  3918. handlers[torch.Size] = lambda tx, value: SizeVariable(
  3919. [create(tx, x) for x in value]
  3920. )
  3921. handlers[collections.OrderedDict] = handlers[dict]
  3922. handlers[immutable_dict] = handlers[dict]
  3923. handlers[immutable_list] = handlers[list]
  3924. # Sourceless MappingProxyType object can be encountered while tracing
  3925. # type.__dict__["__dict__"].__get__
  3926. handlers[types.MappingProxyType] = lambda tx, value: MappingProxyVariable(
  3927. ConstDictVariable(
  3928. {create(tx, k): create(tx, v) for k, v in value.items()},
  3929. dict,
  3930. mutation_type=ValueMutationNew(),
  3931. ),
  3932. )
  3933. handlers[types.GetSetDescriptorType] = (
  3934. lambda tx, value: GetSetDescriptorVariable(value)
  3935. )
  3936. handlers[inspect.Parameter] = lambda tx, value: UserDefinedObjectVariable(
  3937. value, mutation_type=ValueMutationNew()
  3938. )
  3939. handlers[random.Random] = lambda tx, value: RandomClassVariable()
  3940. handlers[types.ModuleType] = lambda tx, value: PythonModuleVariable(value)
  3941. handlers[torch.DispatchKeySet] = lambda tx, value: DispatchKeySetVariable(
  3942. value, mutation_type=ValueMutationNew()
  3943. )
  3944. handlers[torch._functorch.pyfunctorch.FuncTorchInterpreter] = (
  3945. lambda tx, value: FuncTorchInterpreterVariable(
  3946. value, mutation_type=ValueMutationNew()
  3947. )
  3948. )
  3949. handlers[torch.distributions.constraints._Real] = (
  3950. lambda tx, value: UserDefinedObjectVariable(
  3951. value, mutation_type=ValueMutationNew()
  3952. )
  3953. )
  3954. handlers[torch.distributions.constraints._Interval] = (
  3955. lambda tx, value: UserDefinedObjectVariable(
  3956. value, mutation_type=ValueMutationNew()
  3957. )
  3958. )
  3959. handlers[torch.distributions.constraints.Constraint] = (
  3960. lambda tx, value: UserDefinedObjectVariable(
  3961. value, mutation_type=ValueMutationNew()
  3962. )
  3963. )
  3964. def passthrough(tx: "InstructionTranslator", value: T) -> T:
  3965. return value
  3966. for cls in VariableTrackerMeta.all_subclasses:
  3967. handlers[cls] = passthrough
  3968. return handlers
  3969. SourcelessBuilder._type_handlers = SourcelessBuilder.make_type_handlers()
  3970. class SourcelessUserDefinedObjectBuilder:
  3971. """
  3972. SourceLessBuilder does not return a UserDefinedObjectVariable, but in some
  3973. cases it might be ok to return UserDefinedObjects. In such case, use this
  3974. builder.
  3975. """
  3976. def __init__(self) -> None:
  3977. raise AssertionError("Use SourcelessUserDefinedObjectBuilder.create()")
  3978. @staticmethod
  3979. def create(tx: "InstructionTranslator", value: Any) -> VariableTracker:
  3980. value_type = type(value)
  3981. if issubclass(value_type, MutableMapping):
  3982. return MutableMappingVariable(value, mutation_type=ValueMutationNew())
  3983. elif isinstance(value, torch.nn.Module):
  3984. return UnspecializedNNModuleVariable(
  3985. value, mutation_type=ValueMutationNew()
  3986. )
  3987. else:
  3988. return UserDefinedObjectVariable(value, mutation_type=ValueMutationNew())