proxy.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823
  1. # mypy: ignore-errors
  2. import collections
  3. import copy
  4. import dis
  5. import enum
  6. import inspect
  7. import logging
  8. import operator
  9. import sys
  10. import traceback
  11. from collections import OrderedDict
  12. from collections.abc import Iterator
  13. from dataclasses import fields, is_dataclass
  14. from typing import Any, Callable, Optional
  15. import torch
  16. import torch.fx.traceback as fx_traceback
  17. from torch._C import _fx_map_aggregate as map_aggregate, _fx_map_arg as map_arg
  18. from torch.utils._traceback import CapturedTraceback
  19. from ._compatibility import compatibility
  20. from .graph import Graph, magic_methods, reflectable_magic_methods
  21. from .immutable_collections import immutable_dict, immutable_list
  22. from .node import Argument, base_types, Node, Target
  23. from .operator_schemas import check_for_mutable_operation
  24. __all__ = [
  25. "TracerBase",
  26. "GraphAppendingTracer",
  27. "TraceError",
  28. "Proxy",
  29. "MetaProxy",
  30. "Attribute",
  31. "ParameterProxy",
  32. "Scope",
  33. "ScopeContextManager",
  34. ]
  35. log = logging.getLogger(__name__)
  36. @compatibility(is_backward_compatible=False)
  37. class Scope:
  38. """Scope object that records the module path and the module type
  39. of a module. Scope is used to track the information of the module
  40. that contains a Node in a Graph of GraphModule. For example::
  41. class Sub(torch.nn.Module):
  42. def forward(self, x):
  43. # This will be a call_method Node in GraphModule,
  44. # scope for this would be (module_path="sub", module_type=Sub)
  45. return x.transpose(1, 2)
  46. class M(torch.nn.Module):
  47. def __init__(self) -> None:
  48. self.sub = Sub()
  49. def forward(self, x):
  50. # This will be a call_method Node as well,
  51. # scope for this would be (module_path="", None)
  52. x = x.transpose(1, 2)
  53. x = self.sub(x)
  54. return x
  55. """
  56. def __init__(self, module_path: str, module_type: Any):
  57. super().__init__()
  58. self.module_path = module_path
  59. self.module_type = module_type
  60. @compatibility(is_backward_compatible=False)
  61. class ScopeContextManager:
  62. """A context manager to track the Scope of Node during symbolic tracing.
  63. When entering a forward function of a Module, we'll update the scope information of
  64. the current module, and when we exit, we'll restore the previous scope information.
  65. """
  66. def __init__(
  67. self,
  68. scope: Scope,
  69. current_scope: Scope,
  70. ):
  71. super().__init__()
  72. # Keep a copy of prev scope to restore on exit
  73. self._prev_scope = copy.copy(scope)
  74. # Update scope to current scope
  75. scope.module_path = current_scope.module_path
  76. scope.module_type = current_scope.module_type
  77. # Save a reference so we can restore it
  78. self._scope = scope
  79. def __enter__(self):
  80. return self._scope
  81. def __exit__(self, *args):
  82. self._scope.module_path = self._prev_scope.module_path
  83. self._scope.module_type = self._prev_scope.module_type
  84. return
  85. _COPY_META_FIELDS = [
  86. "nn_module_stack",
  87. "torch_fn",
  88. "source_fn_stack",
  89. "original_aten",
  90. "recompute",
  91. "ac_graph_id",
  92. "has_backward_hook",
  93. "from_node",
  94. "quantization_tag", # TODO deprecated
  95. "_numeric_debug_handle", # TODO deprecated
  96. "custom",
  97. "partitioner_tag",
  98. ]
  99. @compatibility(is_backward_compatible=True)
  100. class TracerBase:
  101. graph: Graph
  102. record_stack_traces: bool = False
  103. # When record_stack_traces is True, only reocrd stack traces
  104. # with forward function names.
  105. # This helps when we want stack trace back to model code
  106. _record_forward_stack_traces_only: bool = False
  107. # Feature flag for mutable schema checking
  108. # Enableby default in 1.12
  109. check_mutable_operations: bool = False
  110. # Feature flag for assert tracing
  111. trace_asserts: bool = False
  112. # Feature flag for proxying accesses to buffer values
  113. proxy_buffer_attributes: bool = False
  114. # Name of the function to be traced. It will only be used when
  115. # ``root`` is an instance of ``nn.Module``
  116. traced_func_name: str = "forward"
  117. # Maps the containing module's name to the operator name
  118. scope: Scope
  119. # Records the module call stack
  120. module_stack: OrderedDict[str, tuple[str, Any]]
  121. # Mapping of node name to module scope
  122. node_name_to_scope: dict[str, tuple[str, type]]
  123. @compatibility(is_backward_compatible=True)
  124. def create_node(
  125. self,
  126. kind: str,
  127. target: Target,
  128. args: tuple[Argument, ...],
  129. kwargs: dict[str, Argument],
  130. name: Optional[str] = None,
  131. type_expr: Optional[Any] = None,
  132. ) -> Node:
  133. """
  134. Inserts a graph node given target, args, kwargs, and name.
  135. This method can be overridden to do extra checking, validation, or
  136. modification of values used in node creation. For example, one might
  137. want to disallow in-place operations from being recorded.
  138. """
  139. if kind == "call_function" and self.check_mutable_operations:
  140. check_for_mutable_operation(target, args, kwargs)
  141. node = self.graph.create_node(kind, target, args, kwargs, name, type_expr)
  142. # TODO node_name_to_scope will be depreciated in favor of
  143. # node.meta['nn_module_stack']
  144. self.node_name_to_scope[node.name] = (
  145. self.scope.module_path,
  146. self.scope.module_type,
  147. )
  148. # Optionally set stack trace on the created Node for debugging purposes
  149. if fx_traceback.has_preserved_node_meta():
  150. current_meta: dict[str, Any] = fx_traceback.get_current_meta()
  151. stack_trace = current_meta.get("stack_trace")
  152. if stack_trace:
  153. node.stack_trace = stack_trace
  154. # Explicitly set the stack_trace, nn_module_stack and source_fn on the node.meta
  155. # If other meta fields are needed, they can be added here
  156. for field in _COPY_META_FIELDS:
  157. if field in current_meta:
  158. node.meta[field] = copy.copy(current_meta[field])
  159. # Here we decrement to account for the sequence_nr having
  160. # just been incremented while tracing this lowered aten op.
  161. new_seq_nr = torch.autograd._get_sequence_nr() - 1
  162. # The sequence_nr increments every time a new autograd Node
  163. # is created. During the FWD pass we store the sequence_nr
  164. # corresponding to the last autograd Node created on this fx
  165. # node's meta. A single aten op can create multiple autograd
  166. # nodes as is the case with in-place foreach ops. During the
  167. # BWD pass we retrieve the sequence_nr stored on the current
  168. # executing autograd Node. See NOTE [ Sequence Number ].
  169. if current_meta.get("in_grad_fn", 0) > 0:
  170. new_seq_nr = current_meta["grad_fn_seq_nr"][-1]
  171. node.meta["seq_nr"] = new_seq_nr
  172. elif self.module_stack:
  173. node.meta["nn_module_stack"] = copy.copy(self.module_stack)
  174. if self.record_stack_traces and not node.stack_trace:
  175. user_stack_summary = CapturedTraceback.extract().summary()
  176. if user_stack_summary:
  177. user_stack_summary = self._filter_traceback_frames(user_stack_summary)
  178. if user_stack_summary:
  179. node.stack_trace = "".join(user_stack_summary.format()).strip()
  180. log.debug("create_node %s", node)
  181. return node
  182. def _filter_traceback_frames(
  183. self, user_stack_summary: traceback.StackSummary
  184. ) -> traceback.StackSummary:
  185. # This method can be overridden to customize the frame filtering logic
  186. # for the recorded stack trace
  187. user_frames = []
  188. if self._record_forward_stack_traces_only:
  189. user_frames = [
  190. frame
  191. for frame in user_stack_summary
  192. if (
  193. frame.name == "forward"
  194. or frame.filename.endswith("torch/__init__.py")
  195. )
  196. ]
  197. else:
  198. first_forward = -1
  199. for i, frame in enumerate(user_stack_summary):
  200. if frame.name == "forward":
  201. user_frames = user_stack_summary[i:]
  202. first_forward = i
  203. break
  204. # Not having a "forward" call in the stacktrace implies the
  205. # stacktrace will probably be irrelevant
  206. if first_forward == -1:
  207. user_frames = []
  208. from torch.fx.experimental.symbolic_shapes import uninteresting_files
  209. user_frames = [
  210. frame
  211. for frame in user_frames
  212. if frame.filename not in uninteresting_files()
  213. ]
  214. return traceback.StackSummary.from_list(user_frames)
  215. @compatibility(is_backward_compatible=True)
  216. def proxy(self, node: Node) -> "Proxy":
  217. return Proxy(node, self)
  218. @compatibility(is_backward_compatible=True)
  219. def create_proxy(
  220. self,
  221. kind: str,
  222. target: Target,
  223. args: tuple[Any, ...],
  224. kwargs: dict[str, Any],
  225. name: Optional[str] = None,
  226. type_expr: Optional[Any] = None,
  227. # fix noqa when updating bc tests
  228. proxy_factory_fn: Callable[[Node], "Proxy"] = None, # noqa: RUF013
  229. ):
  230. """
  231. Create a Node from the given arguments, then return the Node
  232. wrapped in a Proxy object.
  233. If kind = 'placeholder', then we're creating a Node that
  234. represents the parameter of a function. If we need to encode
  235. a default parameter, we use the ``args`` tuple. ``args`` is
  236. otherwise empty for ``placeholder`` Nodes.
  237. """
  238. args_ = self.create_arg(args)
  239. kwargs_ = self.create_arg(kwargs)
  240. assert isinstance(args_, tuple)
  241. assert isinstance(kwargs_, dict)
  242. node = self.create_node(kind, target, args_, kwargs_, name, type_expr)
  243. if not proxy_factory_fn:
  244. proxy = self.proxy(node)
  245. else:
  246. proxy = proxy_factory_fn(node)
  247. return proxy
  248. def _find_user_frame(self):
  249. """
  250. Find the Python stack frame executing the user code during
  251. symbolic tracing.
  252. """
  253. # We have to do a little dance here. Basically, walk up the callstack and
  254. # record the first frame not in the pytorch source. This is the frame executing
  255. # the user code during tracing.
  256. frame = inspect.currentframe()
  257. pt_files = [
  258. "torch/fx/proxy.py",
  259. "torch/fx/_symbolic_trace.py",
  260. "torch/fx/experimental/proxy_tensor.py",
  261. "torch/_ops.py",
  262. "torch/_tensor.py",
  263. "torch/utils/_python_dispatch.py",
  264. "torch/_prims_common/wrappers.py",
  265. "torch/_refs/__init__.py",
  266. "torch/_refs/nn/functional/__init__.py",
  267. "torch/utils/_stats.py",
  268. ]
  269. while frame:
  270. frame = frame.f_back
  271. if frame and all(
  272. not frame.f_code.co_filename.endswith(file) for file in pt_files
  273. ):
  274. break
  275. if not frame:
  276. return None
  277. return frame
  278. @compatibility(is_backward_compatible=True)
  279. def create_arg(self, a: Any) -> Argument:
  280. """
  281. A method that lowers the objects seen as arguments during symbolic evaluation
  282. into Argument types that can be stored in IR.
  283. Can be override to support more trace-specific types.
  284. """
  285. # IMPORTANT: Are you here because you are trying to proxy a new type into
  286. # the graph? Please Please Please contact someone on the PyTorch Compiler team;
  287. # the considerations are subtle.
  288. #
  289. # 1) When you add a new type, all of the downstream consumers and pass writers
  290. # need to handle the new type. torch.fx is intended to be easy to write
  291. # passes for, so we will push back against new types.
  292. # 2) In torch.compile's IR, there are only specific operations that go
  293. # into the graph. In particular, Tensor operations should go into the graph,
  294. # but non-Tensor operations shouldn't. What that means is that constructors
  295. # for new types *SHOULD NOT* become nodes in the FX graph.
  296. handler = _create_arg_bypass.get(type(a))
  297. if handler is not None:
  298. # this is just a performance optimization and can be removed if needed
  299. # for common types, we have a fast path to avoid isinstance() overhead
  300. # this doesn't remove the checks below since we need to handle subclasses
  301. return handler(self, a)
  302. if isinstance(a, Proxy):
  303. return a.node # most common arg type goes first
  304. elif hasattr(a, "__fx_create_arg__"):
  305. return a.__fx_create_arg__(self)
  306. # aggregates
  307. elif isinstance(a, tuple):
  308. if hasattr(a, "_fields"):
  309. # NamedTuple constructors don't seem to like getting a generator
  310. # expression as an argument to their constructor, so build this
  311. # intermediate tuple and unpack it into the NamedTuple constructor
  312. args = [self.create_arg(elem) for elem in a]
  313. return type(a)(*args) # type: ignore[arg-type]
  314. return type(a)([self.create_arg(elem) for elem in a])
  315. elif isinstance(a, list):
  316. return [self.create_arg(elem) for elem in a]
  317. elif isinstance(a, dict):
  318. return _create_arg_dict(self, a)
  319. elif isinstance(a, slice):
  320. return slice(
  321. self.create_arg(a.start),
  322. self.create_arg(a.stop),
  323. self.create_arg(a.step),
  324. )
  325. elif isinstance(a, range):
  326. return range(
  327. self.create_arg(a.start),
  328. self.create_arg(a.stop),
  329. self.create_arg(a.step),
  330. )
  331. elif isinstance(a, (torch._ops.OpOverload, torch._ops.HigherOrderOperator)):
  332. return a
  333. elif is_dataclass(a):
  334. kwargs = {
  335. field.name: self.create_arg(getattr(a, field.name))
  336. for field in fields(a)
  337. }
  338. return self.create_node("call_function", a.__class__, (), kwargs)
  339. elif isinstance(a, (*base_types, enum.Enum)) or a is None or a is ...:
  340. return a
  341. raise NotImplementedError(f"argument of type: {type(a)}")
  342. @compatibility(is_backward_compatible=True)
  343. def to_bool(self, obj: "Proxy") -> bool:
  344. """Called when a proxy object is being converted to a boolean, such as
  345. when used in control flow. Normally we don't know what to do because
  346. we don't know the value of the proxy, but a custom tracer can attach more
  347. information to the graph node using create_node and can choose to return a value.
  348. """
  349. raise TraceError(
  350. "symbolically traced variables cannot be used as inputs to control flow"
  351. )
  352. @compatibility(is_backward_compatible=True)
  353. def iter(self, obj: "Proxy") -> Iterator:
  354. """Called when a proxy object is being iterated over, such as
  355. when used in control flow. Normally we don't know what to do because
  356. we don't know the value of the proxy, but a custom tracer can attach more
  357. information to the graph node using create_node and can choose to return an iterator.
  358. """
  359. raise TraceError(
  360. "Proxy object cannot be iterated. This can be "
  361. "attempted when the Proxy is used in a loop or"
  362. " as a *args or **kwargs function argument. "
  363. "See the torch.fx docs on pytorch.org for a "
  364. "more detailed explanation of what types of "
  365. "control flow can be traced, and check out the"
  366. " Proxy docstring for help troubleshooting "
  367. "Proxy iteration errors"
  368. )
  369. @compatibility(is_backward_compatible=True)
  370. def keys(self, obj: "Proxy") -> Any:
  371. """Called when a proxy object is has the keys() method called.
  372. This is what happens when ** is called on a proxy. This should return an
  373. iterator it ** is suppose to work in your custom tracer.
  374. """
  375. return Attribute(obj, "keys")()
  376. # used in Proxy object when just appending to the graph while not tracing.
  377. @compatibility(is_backward_compatible=True)
  378. class GraphAppendingTracer(TracerBase):
  379. def __init__(self, graph: Graph):
  380. super().__init__()
  381. self.graph = graph
  382. self.scope = Scope("", None)
  383. self.module_stack = collections.OrderedDict()
  384. self.node_name_to_scope = {}
  385. @compatibility(is_backward_compatible=False)
  386. def assert_fn(x):
  387. assert x
  388. @compatibility(is_backward_compatible=True)
  389. class TraceError(ValueError):
  390. pass
  391. @compatibility(is_backward_compatible=True)
  392. class Proxy:
  393. """
  394. ``Proxy`` objects are ``Node`` wrappers that flow through the
  395. program during symbolic tracing and record all the operations
  396. (``torch`` function calls, method calls, operators) that they touch
  397. into the growing FX Graph.
  398. If you're doing graph transforms, you can wrap your own ``Proxy``
  399. method around a raw ``Node`` so that you can use the overloaded
  400. operators to add additional things to a ``Graph``.
  401. ``Proxy`` objects cannot be iterated. In other words, the symbolic
  402. tracer will throw an error if a ``Proxy`` is used in a loop or as
  403. an ``*args``/``**kwargs`` function argument.
  404. There are two main ways around this:
  405. 1. Factor out the untraceable logic into a top-level function and
  406. use ``fx.wrap`` on it.
  407. 2. If the control flow is static (i.e. the loop trip count is
  408. based on some hyperparameter), the code can be kept in its original
  409. position and refactored into something like::
  410. for i in range(self.some_hyperparameter):
  411. indexed_item = proxied_value[i]
  412. For a more detailed description into the Proxy internals, check out
  413. the "Proxy" section in `torch/fx/README.md`
  414. """
  415. @compatibility(is_backward_compatible=True)
  416. def __init__(self, node: Node, tracer: "Optional[TracerBase]" = None):
  417. if tracer is None:
  418. # This allows you to create a Proxy object around a raw Node
  419. tracer = GraphAppendingTracer(node.graph)
  420. self.tracer = tracer
  421. self.node = node
  422. def __repr__(self) -> str:
  423. return f"Proxy({self.node.name})"
  424. def __getattr__(self, k) -> "Attribute":
  425. # note: not added to the graph yet, if this is a method call
  426. # we peephole optimize to the method invocation
  427. return Attribute(self, k)
  428. def __getstate__(self) -> dict:
  429. return self.__dict__
  430. def __deepcopy__(self, memo) -> dict:
  431. # We have to explicitly override this method, because otherwise deepcopy
  432. # will go to __getattr__(self, "__deepcopy__") and return a
  433. # Attribute(__deepcopy__), and may go into an infinite loop in some cases.
  434. import copy
  435. new_dict = {}
  436. for k, v in self.__dict__.items():
  437. try:
  438. new_obj = copy.deepcopy(v, memo)
  439. except Exception:
  440. log.warning(
  441. "Shallow copy %s of Proxy because it cannot be deepcopied. "
  442. "Proxy is created for node %s",
  443. k,
  444. self.node.name,
  445. )
  446. new_obj = copy.copy(v)
  447. new_dict[k] = new_obj
  448. assert "node" in new_dict
  449. assert "tracer" in new_dict
  450. new_proxy = Proxy(new_dict["node"], new_dict["tracer"])
  451. for k, v in new_dict.items():
  452. new_proxy.__dict__[k] = v
  453. return new_proxy
  454. def __setstate__(self, d):
  455. # This is called when being unpickled/loaded.
  456. self.__dict__ = d
  457. def __call__(self, *args, **kwargs) -> "Proxy":
  458. return self.tracer.create_proxy(
  459. "call_method", "__call__", (self,) + args, kwargs
  460. )
  461. def __iter__(self) -> Iterator["Proxy"]:
  462. frame = inspect.currentframe()
  463. assert frame is not None
  464. calling_frame = frame.f_back
  465. assert calling_frame is not None
  466. inst_list = list(dis.get_instructions(calling_frame.f_code))
  467. if sys.version_info >= (3, 11):
  468. from bisect import bisect_left
  469. inst_idx = bisect_left(
  470. inst_list, calling_frame.f_lasti, key=lambda x: x.offset
  471. )
  472. else:
  473. inst_idx = calling_frame.f_lasti // 2
  474. inst = inst_list[inst_idx]
  475. if inst.opname == "UNPACK_SEQUENCE":
  476. return (self[i] for i in range(inst.argval)) # type: ignore[index]
  477. return self.tracer.iter(self)
  478. def __abs__(self):
  479. return self.tracer.create_proxy("call_function", operator.abs, (self,), {})
  480. def __bool__(self) -> bool:
  481. if self.tracer.trace_asserts:
  482. # check if this boolean is used in an assertion, bytecode pattern for assertions
  483. # is pretty stable for Python 3.7--3.9
  484. frame = inspect.currentframe()
  485. assert frame is not None
  486. calling_frame = frame.f_back
  487. assert calling_frame is not None
  488. insts = list(dis.get_instructions(calling_frame.f_code))
  489. if sys.version_info >= (3, 11):
  490. from bisect import bisect_left
  491. cur = bisect_left(insts, calling_frame.f_lasti, key=lambda x: x.offset)
  492. else:
  493. cur = calling_frame.f_lasti // 2
  494. inst = insts[cur]
  495. if inst.opname == "POP_JUMP_IF_TRUE":
  496. first = insts[cur + 1]
  497. assert inst.arg is not None
  498. last = insts[inst.arg // 2 - 1]
  499. starts_with_assert = (
  500. first.opname == "LOAD_GLOBAL"
  501. and first.argval == "AssertionError"
  502. or first.opname == "LOAD_ASSERTION_ERROR"
  503. )
  504. if starts_with_assert and last.opname == "RAISE_VARARGS":
  505. self.tracer.create_proxy("call_function", assert_fn, (self,), {})
  506. return True
  507. return self.tracer.to_bool(self)
  508. @compatibility(is_backward_compatible=True)
  509. def keys(self):
  510. return self.tracer.keys(self)
  511. def __len__(self):
  512. raise RuntimeError(
  513. "'len' is not supported in symbolic tracing by default. If you want "
  514. "this call to be recorded, please call torch.fx.wrap('len') at "
  515. "module scope"
  516. )
  517. @classmethod
  518. def __torch_function__(cls, orig_method, types, args=None, kwargs=None):
  519. args = args if args else ()
  520. kwargs = kwargs if kwargs else {}
  521. tracers: dict[Any, None] = {}
  522. def find_tracer(a):
  523. if isinstance(a, cls):
  524. tracers[a.tracer] = None
  525. map_aggregate(args, find_tracer)
  526. map_aggregate(kwargs, find_tracer)
  527. if len(tracers) > 1:
  528. raise RuntimeError(
  529. f"Found multiple different tracers {list(tracers.keys())} while "
  530. f"trying to trace operations {orig_method}"
  531. )
  532. tracer = next(iter(tracers.keys()))
  533. if isinstance(orig_method, torch._C.ScriptMethod):
  534. args = (orig_method.owner,) + args
  535. return tracer.create_proxy("call_method", orig_method.name, args, kwargs)
  536. if torch.overrides.is_tensor_method_or_property(orig_method):
  537. return tracer.create_proxy(
  538. "call_method", orig_method.__name__, args, kwargs
  539. )
  540. else:
  541. if isinstance(orig_method, torch._ops.HigherOrderOperator):
  542. # TODO: Define how to symbolically trace HigherOrderOperators
  543. raise RuntimeError("Unable to symbolically trace HigherOrderOperators")
  544. return tracer.create_proxy(
  545. "call_function",
  546. orig_method,
  547. args,
  548. kwargs,
  549. name=tracer.graph._target_to_str(orig_method.__name__),
  550. )
  551. @compatibility(is_backward_compatible=False)
  552. class MetaProxy(Proxy):
  553. """
  554. A Proxy subclass that propagates metadata (meta['val']) during graph tracing.
  555. """
  556. def __init__(
  557. self, node: Node, tracer: "Optional[TracerBase]" = None, fake_mode=None
  558. ):
  559. super().__init__(node, tracer)
  560. self.fake_mode = fake_mode
  561. def __repr__(self) -> str:
  562. return f"MetaProxy({self.node.name})"
  563. @classmethod
  564. def __torch_function__(cls, orig_method, types, args=None, kwargs=None):
  565. args = args if args else ()
  566. kwargs = kwargs if kwargs else {}
  567. meta_proxy = None
  568. for arg in args:
  569. if isinstance(arg, MetaProxy):
  570. meta_proxy = arg
  571. break
  572. assert meta_proxy is not None, (
  573. "No MetaProxy found in arguments, but one is expected."
  574. )
  575. proxy = super().__torch_function__(orig_method, types, args, kwargs)
  576. with meta_proxy.fake_mode:
  577. proxy.node.meta["val"] = orig_method(
  578. *[a.node.meta["val"] if isinstance(a, Proxy) else a for a in args],
  579. **kwargs,
  580. )
  581. return MetaProxy(proxy.node, proxy.tracer, meta_proxy.fake_mode)
  582. @compatibility(is_backward_compatible=True)
  583. class Attribute(Proxy):
  584. @compatibility(is_backward_compatible=True)
  585. def __init__(self, root: Proxy, attr: str):
  586. self.root = root
  587. self.attr = attr
  588. self.tracer = root.tracer
  589. self._node: Optional[Node] = None
  590. @property
  591. def node(self):
  592. # the node for attributes is added lazily, since most will just be method calls
  593. # which do not rely on the getitem call
  594. if self._node is None:
  595. self._node = self.tracer.create_proxy(
  596. "call_function", getattr, (self.root, self.attr), {}
  597. ).node
  598. return self._node
  599. def __call__(self, *args, **kwargs):
  600. return self.tracer.create_proxy(
  601. "call_method", self.attr, (self.root,) + args, kwargs
  602. )
  603. @compatibility(is_backward_compatible=False)
  604. class ParameterProxy(Proxy):
  605. """
  606. A special proxy which lets "shape", "size", "dim", and a few other
  607. attribute accesses pass through to the underlying module parameter object,
  608. so that conditional tests on these attributes will not throw exception during tracing
  609. """
  610. def __init__(self, tracer: TracerBase, node: Node, name, param):
  611. super().__init__(node, tracer)
  612. assert isinstance(param, torch.nn.Parameter)
  613. self.param = param
  614. self.name = name
  615. def __repr__(self) -> str:
  616. return f"ParameterProxy({self.name})"
  617. @property
  618. def shape(self):
  619. return self.param.shape
  620. def size(self):
  621. return self.param.size()
  622. def dim(self):
  623. return self.param.dim()
  624. @property
  625. def ndim(self):
  626. return self.param.ndim
  627. def numel(self):
  628. return self.param.numel()
  629. def nelement(self):
  630. return self.param.nelement()
  631. for method in magic_methods:
  632. def _scope(method):
  633. def impl(*args, **kwargs):
  634. tracer = args[0].tracer
  635. target = getattr(operator, method)
  636. return tracer.create_proxy("call_function", target, args, kwargs)
  637. impl.__name__ = method
  638. as_magic = f"__{method.strip('_')}__"
  639. setattr(Proxy, as_magic, impl)
  640. _scope(method)
  641. def _define_reflectable(orig_method_name):
  642. method_name = f"__r{orig_method_name.strip('_')}__"
  643. def impl(self, rhs):
  644. target = getattr(operator, orig_method_name)
  645. return self.tracer.create_proxy("call_function", target, (rhs, self), {})
  646. impl.__name__ = method_name
  647. impl.__qualname__ = method_name
  648. setattr(Proxy, method_name, impl)
  649. for orig_method_name in reflectable_magic_methods:
  650. _define_reflectable(orig_method_name)
  651. def _no_nodes_error(arg):
  652. raise RuntimeError(
  653. "Keys for dictionaries used as an argument cannot contain a "
  654. f"Node. Got key: {arg}"
  655. )
  656. def _create_arg_dict(self, a):
  657. r = {}
  658. for k, v in a.items():
  659. if not isinstance(k, str):
  660. # Check for invalid dict keys. We do not want a Proxy to appear
  661. # anywhere within the key. Since keys can be collection types,
  662. # we iterate through the key with map_arg
  663. k = self.create_arg(k)
  664. map_arg(k, _no_nodes_error)
  665. r[k] = self.create_arg(v)
  666. return r
  667. _create_arg_bypass = {
  668. t: lambda self, a: a
  669. for t in [
  670. *base_types,
  671. type(None),
  672. type(...),
  673. torch._ops.OpOverload,
  674. torch._ops.HigherOrderOperator,
  675. ]
  676. }
  677. _create_arg_bypass[Proxy] = lambda self, a: a.node
  678. _create_arg_bypass[tuple] = lambda self, a: tuple([self.create_arg(elem) for elem in a])
  679. _create_arg_bypass[list] = lambda self, a: [self.create_arg(elem) for elem in a]
  680. _create_arg_bypass[dict] = _create_arg_dict
  681. _create_arg_bypass[immutable_list] = _create_arg_bypass[list]
  682. _create_arg_bypass[immutable_dict] = _create_arg_bypass[dict]