""" This module provides decorators and utilities for controlling TorchDynamo's behavior during compilation. """ import functools import inspect import weakref from collections.abc import Callable from dataclasses import dataclass from types import TracebackType from typing import Any, overload, TYPE_CHECKING, TypeVar from typing_extensions import ParamSpec import torch import torch.utils._pytree as pytree from torch._opaque_base import OpaqueBase from torch._vendor.packaging.version import InvalidVersion, Version from torch.compiler import is_compiling from torch.utils._contextlib import _DecoratorContextManager from torch.utils._python_dispatch import is_traceable_wrapper_subclass from .._utils_internal import justknobs_check from . import trace_rules, variables from .comptime import comptime from .eval_frame import ( _set_stance, DisableContext, DynamoStance, innermost_fn, RunOnlyContext, skip_code, ) from .external_utils import ( get_nonrecursive_disable_wrapper, wrap_dunder_call_ctx_manager, ) from .utils import ( _get_error_on_graph_break, _set_error_on_graph_break, allow_lru_cache_wrapper_trace_without_warning, is_function, is_lru_cache_wrapped_function, ) justknobs_check._dynamo_marked_constant = True # type: ignore[attr-defined] if TYPE_CHECKING: from types import FunctionType from torch._C._dynamo.eval_frame import ( # noqa: F401 reset_code, set_eval_frame, set_guard_complete_hook, set_guard_error_hook, unsupported, ) from .variables import VariableTracker else: for name in dir(torch._C._dynamo.eval_frame): if name.startswith("__"): continue globals()[name] = getattr(torch._C._dynamo.eval_frame, name) _P = ParamSpec("_P") _R = TypeVar("_R") FuncType = Callable[..., Any] F = TypeVar("F", bound=FuncType) def run(fn: Callable[_P, _R] | None = None) -> Any: """Don't do any dynamic compiles, just use prior optimizations""" if fn is not None: fn = innermost_fn(fn) if not callable(fn): raise AssertionError("fn must be callable") return RunOnlyContext()(fn) return RunOnlyContext() def disable(fn=None, recursive=True, *, reason=None, wrapping=True): # type: ignore[no-untyped-def] """ Decorator to disable TorchDynamo If recursive=True, Dynamo is completely skipped on the decorated function frame as well as the recursively invoked functions. If recursive=False, Dynamo skips frames associated with the function code, but still process recursively invoked frames. If reason is provided, it will be printed when Dynamo attempts to trace the disabled function. """ if recursive: if fn is not None: fn = innermost_fn(fn) if not callable(fn): raise AssertionError("fn must be callable") return DisableContext(msg=reason, wrapping=wrapping)(fn) return DisableContext(msg=reason, wrapping=wrapping) else: def wrap(fn: Callable[_P, _R]) -> Callable[_P, _R]: fn = innermost_fn(fn) if not callable(fn): raise AssertionError("fn must be callable") nonrecursive_disable_wrapper = get_nonrecursive_disable_wrapper(fn) nonrecursive_disable_wrapper._torchdynamo_disable = True # type: ignore[attr-defined] nonrecursive_disable_wrapper._torchdynamo_disable_msg = reason # type: ignore[attr-defined] nonrecursive_disable_wrapper._torchdynamo_orig_callable = fn # type: ignore[attr-defined] nonrecursive_disable_wrapper._torchdynamo_wrapper_id = id( # type: ignore[attr-defined] nonrecursive_disable_wrapper ) nonrecursive_disable_wrapper._torchdynamo_disable_recursive = False # type: ignore[attr-defined] return nonrecursive_disable_wrapper if fn is None: return wrap return wrap(fn) _nonrecursive_disable_wrapper_code = disable(lambda: None, recursive=False).__code__ # type: ignore[attr-defined] skip_code(_nonrecursive_disable_wrapper_code) def skip(fn: Callable[_P, _R] | None = None) -> Callable[..., Any]: """ Skip frames associated with the function code, but still process recursively invoked frames """ if fn is None: return skip fn = innermost_fn(fn) if not callable(fn): raise AssertionError("fn must be callable") skip_code(fn.__code__) fn._torchdynamo_disable = True # type: ignore[attr-defined] return fn class set_stance(_DecoratorContextManager): """ Decorator, context manager, function to set the current stance of the compiler. Stances documented in corresponding function in torch/compiler/__init__.py """ _dynamo_forbidden = True def __init__( self, stance: str = "default", *, skip_guard_eval_unsafe: bool = False, force_backend: str | Callable[..., Any] | None = None, ) -> None: if force_backend is not None and stance != "default": raise RuntimeError("non-default stance cannot have force_backend set") self.stance = DynamoStance(stance, skip_guard_eval_unsafe, force_backend) self.prev = _set_stance(self.stance) def __call__(self, fn: F) -> F: _set_stance(self.prev) wrapper = super().__call__(fn) # forbid wrapper in graph wrapper._dynamo_forbidden = True # type: ignore[attr-defined] return wrapper def __enter__(self) -> None: _set_stance(self.stance) def __exit__( self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None, ) -> None: _set_stance(self.prev) def clone(self) -> "set_stance": return self.__class__(self.stance.stance, force_backend=self.stance.backend) def assume_constant_result(fn): # type: ignore[no-untyped-def] fn._dynamo_marked_constant = True # type: ignore[attr-defined] return fn def allow_in_graph(fn): # type: ignore[no-untyped-def] """ Tells the compiler frontend (Dynamo) to skip symbolic introspection of the function and instead directly write it to the graph when encountered. See :func:`torch.compiler.allow_in_graph`'s docstring for the full documentation WARNING: this API can be a footgun, please read the documentation carefully. """ if isinstance(fn, (list, tuple)): return [allow_in_graph(x) for x in fn] if not callable(fn): raise AssertionError("allow_in_graph expects a callable") if trace_rules.lookup_callable(fn) != variables.TorchInGraphFunctionVariable: fn_id = id(fn) trace_rules._disallowed_callable_ids.remove(fn_id) trace_rules._allowed_callable_ids.add(fn_id) # Avoid id reuse which creates subtle bugs. def deregister() -> None: trace_rules._allowed_callable_ids.remove(fn_id) weakref.finalize(fn, deregister) return fn # pyrefly: ignore [implicit-any] def _check_mutually_exclusive_decorators(fn: Callable, decorator_name: str) -> None: mutually_exclusive = { "leaf_function": trace_rules.is_leaf_function, "nonstrict_trace": trace_rules.is_nonstrict_trace_callable, } for other_name, check_fn in mutually_exclusive.items(): if other_name != decorator_name and check_fn(fn): first, second = sorted([decorator_name, other_name]) raise ValueError( f"Function {fn} cannot be both marked as @{first} and " f"@{second}. Please use only one decorator." ) def nonstrict_trace(traceable_fn: Callable[_P, _R]) -> Callable[_P, _R]: """ Decorator to mark a function as nonstrict-traceable for dynamo. A nonstrict-traced function appears as an opaque call in the dynamo graph. Dynamo does not trace into the function body (hence the "nonstrict"), but aot_autograd will trace into it. This is similar to ``allow_in_graph`` but with enhanced support for: - User-defined classes as inputs (must be registered with pytree) - ``nn.Module`` as input arguments (parameters and buffers are tracked for autograd) - Global/captured tensors treated as constants (assumed not updated during execution) Note: - With ``backend="eager"``, the original Python function runs directly. With ``backend="aot_eager"``, the graph traced by aot_autograd runs. With ``backend="inductor"``, the traced graph is compiled with inductor. - Training is supported: you can call ``.backward()`` on outputs and gradients will flow through the nonstrict-traced function. Dangerous patterns (may cause silent incorrectness): - Side effects between nonstric_traced fn and compiled region: The function should not depend on variables mutated by other code inside the compiled function, and code after the call should not depend on mutations made by it. - Implicit inputs (closures/globals): Tensors captured from enclosing scopes are treated as constants. Gradients will NOT flow back to them. Pass tensors as explicit arguments if gradients are needed. Restrictions: - Both inputs and outputs must use pytree-compatible types. User-defined classes must be registered via :func:`torch.utils._pytree.register_pytree_node`, :func:`torch.utils._pytree.register_dataclass`, or :func:`torch.utils._pytree.register_constant`. Tensors, Python primitives (int, float, bool, str), symbolic types (SymInt, SymFloat, SymBool), and built-in containers (list, tuple, dict) are already handled by default. - Primitive values and container structure are specialized per call site: each call site expects the same primitives and structure on every execution. Example:: >>> import torch >>> @torch._dynamo.nonstrict_trace ... def traced_forward(model, x): ... # It's OK to have dynamo graph break within nonstrict_trace region ... torch._dynamo.graph_break() ... return model(x) + x ... >>> class MyModule(torch.nn.Module): ... def __init__(self): ... super().__init__() ... self.inner = torch.nn.Linear(10, 10) ... ... def forward(self, x): ... return traced_forward(self.inner, x) ... >>> # Compile and run >>> model = MyModule() >>> opt_model = torch.compile(model, backend="aot_eager", fullgraph=True) >>> out = opt_model(torch.randn(10, 10)) >>> out.sum().backward() # Gradients flow through traced_forward """ if not callable(traceable_fn): raise AssertionError("nonstrict_trace expects a callable") _check_mutually_exclusive_decorators(traceable_fn, "nonstrict_trace") @functools.wraps(traceable_fn) def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R: return traceable_fn(*args, **kwargs) wrapped_id = id(wrapped) # This line allows us to reuse much of the `allow_in_graph` impl. trace_rules._allowed_callable_ids.add(wrapped_id) # This line allows us to diverge the impl from `allow_in_graph`. trace_rules._nonstrict_trace_callable_ids.add(wrapped_id) # Avoid id reuse which creates subtle bugs. def deregister() -> None: trace_rules._allowed_callable_ids.remove(wrapped_id) trace_rules._nonstrict_trace_callable_ids.remove(wrapped_id) weakref.finalize(wrapped, deregister) return wrapped def _invoke_leaf_function_python( real_impl: Callable[..., Any], fake_impl: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any], mutates_args: frozenset[str] | None = None, hook_fn: Callable[..., Any] | None = None, hook_fake_fn: Callable[..., Any] | None = None, ) -> Any: """Call invoke_leaf_function HOP directly from Python. This enables @leaf_function to work with make_fx without relying on Dynamo to intercept the call. """ from torch._higher_order_ops.invoke_leaf_function import ( _LeafCallable, convert_modules_to_states, invoke_leaf_function, make_leaf_function_wrappers, store_makefx_modules, ) captured_modules: list[torch.nn.Module] = [] seen_module_ids: dict[int, int] = {} # id(module) -> position in captured_modules for val in pytree.tree_flatten( (args, kwargs), is_leaf=lambda x: isinstance(x, torch.nn.Module) )[0]: if isinstance(val, torch.nn.Module) and id(val) not in seen_module_ids: seen_module_ids[id(val)] = len(captured_modules) captured_modules.append(val) global_indices = store_makefx_modules(captured_modules) module_to_index = { mod_id: global_indices[pos] for mod_id, pos in seen_module_ids.items() } processed = convert_modules_to_states((args, kwargs), module_to_index) flat_args, input_spec = pytree.tree_flatten(processed) # Single-element mutable list so the wrappers can write back the output # TreeSpec. Read captured_out_spec[0] after the wrappers have been called. captured_out_spec: list[pytree.TreeSpec | None] = [None] wrapped_real, wrapped_fake = make_leaf_function_wrappers( real_impl, fake_impl, captured_out_spec ) real_fn_callable = _LeafCallable(wrapped_real) fake_fn_callable = _LeafCallable(wrapped_fake) if hook_fn is not None: real_fn_callable._leaf_hook_real_fn = hook_fn # type: ignore[attr-defined] real_fn_callable._leaf_hook_fake_fn = hook_fake_fn # type: ignore[attr-defined] mutated_flat_indices = "" if mutates_args: from torch._higher_order_ops.invoke_leaf_function import ( _resolve_mutated_flat_indices, ) mutated_flat_indices = _resolve_mutated_flat_indices( real_impl, mutates_args, len(flat_args), input_spec ) flat_out = invoke_leaf_function( real_fn_callable, fake_fn_callable, input_spec, mutated_flat_indices, *flat_args ) if captured_out_spec[0] is None: raise AssertionError("captured_out_spec was not set by leaf function wrappers") return pytree.tree_unflatten(flat_out, captured_out_spec[0]) @overload def leaf_function(fn: Callable[_P, _R]) -> Callable[_P, _R]: ... @overload def leaf_function( *, mutates_args: set[str] ) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: ... def leaf_function( fn: Callable[_P, _R] | None = None, *, mutates_args: set[str] | None = None ) -> Callable[_P, _R] | Callable[[Callable[_P, _R]], Callable[_P, _R]]: """ Decorator to mark a function as a leaf function for :func:`torch.compile`. A leaf function appears as an opaque operation in the compiled graph. During compilation, Dynamo and AOT autograd do not trace into it. At runtime, the original eager Python code is executed directly. Quick Start: Suppose we want to log per-sample statistics during the forward pass:: def compute_and_log_stats(x): stats = x.mean(dim=1) logger.info(f"Per-sample means: {stats}") return (stats,) With ``torch.compile(..., fullgraph=True)``, this fails because ``logger.info`` causes a graph break. To fix it, follow the steps below: 1. Decorate your function with ``@leaf_function`` 2. Define a shape-inference function using ``@your_fn.register_fake`` Concrete implementation:: >>> import logging >>> import torch >>> from torch._dynamo.decorators import leaf_function >>> >>> logging.basicConfig(level=logging.INFO) >>> logger = logging.getLogger(__name__) >>> >>> @leaf_function ... def compute_and_log_stats(x): ... stats = x.mean(dim=1) # Shape: (x.shape[0],) ... logger.info(f"Per-sample means: {stats}") ... return (stats,) ... >>> @compute_and_log_stats.register_fake ... def compute_and_log_stats_fake(x): ... # Match the output shape: (x.shape[0],) ... return (x.new_empty(x.shape[0]),) ... >>> def fn(x): ... return compute_and_log_stats(x) ... >>> x = torch.randn(3, 4) >>> compiled_fn = torch.compile(fn, backend="aot_eager", fullgraph=True) >>> out = compiled_fn(x)[0] When to Use: Use ``leaf_function`` when you need eager execution semantics that tracing cannot preserve: - **Non-traceable code**: The function consists of code that Dynamo or AOT Autograd cannot trace. - **Runtime side effects**: Operations that must happen exactly once per call at runtime (e.g., logging, external library calls). If your function has a static computation graph and no runtime side effects, prefer :func:`allow_in_graph` or :func:`nonstrict_trace` instead. They allow AOT autograd to trace through and potentially optimize the code. Usage: **Inputs and Outputs**: - Both inputs and outputs must use pytree-compatible types. Tensors, Python primitives (int, float, bool, str), and built-in containers (list, tuple, dict) are supported by default. Note: We recommend leaf_functions only accept and return tensors. Though primitive types (int, float, bool, str) are supported in inputs and outputs, they may cause surprising behavior: primitive inputs are guarded and may trigger recompilation, while primitive output values are captured from the fake implementation at compile time and remain fixed for all subsequent executions. Even though the real function runs at runtime, its primitive return values are silently replaced with the compile-time values from the fake implementation. Example:: # BAD: primitive output varies at runtime counter = 0 @leaf_function def count_calls(x): global counter counter += 1 return (x, counter) @count_calls.register_fake def count_calls_fake(x): return (x, 999) # placeholder value # At runtime: real function runs (counter increments to 1, 2, 3, ...) # But returned count is always 999 (from fake implementation at compile time) - User-defined classes must be registered via :func:`torch.utils._pytree.register_pytree_node` or :func:`torch.utils._pytree.register_dataclass`. - :class:`torch.nn.Module` can also be passed as input; its parameters and buffers are tracked for autograd. The module must exist outside the compile region. Running the module multiple times must not modify its parameters, buffers, or attributes. **register_fake (required)**: Since the function body is not traced, you must provide a shape-inference function via ``@fn.register_fake``. It runs at compile time with FakeTensor inputs (tensors with no data, only metadata) and must satisfy the following requirements: - Must have the same input and output signature (e.g., same pytree structure, same tensor metadata such as shapes, dtypes, device, strides) as the real function - Must be runnable with FakeTensor inputs - Must only use its explicit arguments (no closures over tensors or modules) Note: The input and output signature must be determinable at compile time. If your function's output structure depends on runtime values, adjust the leaf function boundary until outputs become predictable. To validate that your fake implementation matches the real function's outputs, set ``torch._dynamo.config.leaf_function_validate_outputs = True``. **register_multi_grad_hook (optional)**: You can register a backward hook via ``@fn.register_multi_grad_hook`` to run code when gradients have been computed for all requires_grad tensor inputs during backward. The hook fires exactly once per backward pass. The hook receives only the gradients of the leaf function's requires_grad tensor inputs, in the order they appeared; non-tensor arguments and tensors without ``requires_grad`` are not forwarded to the hook. The hook must return ``None``. The hook is called as a leaf function itself, so it is also opaque to the compiler. Example:: >>> @leaf_function ... def debug_log(t): ... return None ... >>> @debug_log.register_fake ... def debug_log_fake(t): ... return None ... >>> @debug_log.register_multi_grad_hook ... def debug_log_hook(t_grad): ... print(f"[bwd] norm={t_grad.norm().item()}") ... >>> x = torch.randn(4, requires_grad=True) >>> debug_log(x) # no assignment needed >>> (x * 2).sum().backward() [bwd] norm=... Limitations: Currently, inductor backend and :func:`torch.export.export` are not yet supported. Training / Autograd: Training is supported automatically if your leaf function is differentiable in eager mode (e.g., it's implemented with PyTorch ops, :class:`torch.autograd.Function`, or differentiable custom ops). **Restriction**: Calling ``.backward()`` *inside* the leaf function is not supported. **Escaped gradients check**: If the leaf function closes over a tensor with ``requires_grad=True``, gradients will not flow back to it (only explicit inputs receive gradients). To detect such cases, set ``torch._dynamo.config.leaf_function_check_escaped_gradients = True``. When enabled, a ``RuntimeError`` is raised with details about the escaped tensors. Internally, inputs and outputs are detached from the outer autograd graph at the leaf function boundary. The leaf function builds its own local autograd graph. In backward, gradients propagate through this local graph to the leaf function's inputs. If a :class:`torch.nn.Module` is passed in as input, the gradients will also flow back to its parameters and buffers. How leaf_function Works: Understanding this helps avoid pitfalls: 1. **Compilation**: Dynamo and AOT autograd do not trace into the leaf function. Only your fake implementation runs during compilation to determine output signatures, shapes, and dtypes. The real function runs at runtime as eager Python. 2. **Isolation**: Mutations to shared state (globals, closures) may not be visible across the boundary between leaf functions and compiled regions. Pass data explicitly as function arguments and return results as outputs. Dangerous Patterns: These patterns may cause silent incorrectness or errors: - **Side effects between leaf functions and compiled regions**: Mutations to Python state inside a leaf function may not be visible to the compiled region, and vice versa. The compiled graph may reorder or eliminate operations in ways that break assumptions about side effect ordering. Bad:: # Compiled region self.counter += 1 y = my_leaf_fn(self, x) # Don't depend on self.counter inside leaf_fn y = my_leaf_fn(self, x) result = self.state # Don't depend on state mutated by leaf_fn Logging/printing inside the leaf function is fine since it doesn't affect correctness. - **In-place mutations on inputs**: Undeclared in-place mutations on input tensors are detected and will raise an error. Either declare mutations via ``mutates_args`` or clone inputs before mutating. Bad:: @leaf_function def my_leaf_fn(x): x.add_(1) # Will raise: "Undeclared in-place mutation" return (x,) Good:: @leaf_function(mutates_args={"buf"}) def my_leaf_fn(x, buf): buf.add_(1) # OK: declared in mutates_args return (x + buf,) For pytree args (lists, tuples, dicts of tensors), use the parameter name to declare mutation on all contained tensors:: @leaf_function(mutates_args={"buffers"}) def my_leaf_fn(x, buffers): for buf in buffers: buf.add_(1) # OK: 'buffers' declared in mutates_args return (x + sum(buffers),) Or use bracket notation for fine-grained control:: @leaf_function(mutates_args={"buffers[0]"}) def my_leaf_fn(x, buffers): buffers[0].add_(1) # OK: 'buffers[0]' declared return (x + buffers[1],) # buffers[1] is not cloned - **Closures in fake implementation**: Tensors or modules captured from enclosing scopes in the fake implementation will cause compilation errors. The real function can close over them if they don't require gradient, but the fake implementation must only use its arguments. Bad:: weight = torch.randn(3, 3) @leaf_function def my_leaf_fn(x): return (x @ weight,) @my_leaf_fn.register_fake def my_leaf_fn_fake(x): return (x @ weight,) # Error: weight used in fake implementation Good:: weight = torch.randn(3, 3) @leaf_function def my_leaf_fn(x): return (x @ weight,) # OK: real function can use closure @my_leaf_fn.register_fake def my_leaf_fn_fake(x): return (x @ torch.empty_like(x),) # OK: uses only args Example:: Wrapping a linear function with logging as a leaf_function. Gradients flow back because the inner operations are differentiable:: >>> import torch >>> import torch.nn.functional as F >>> from torch._dynamo.decorators import leaf_function >>> >>> @leaf_function ... def custom_forward(linear, x): ... # Logging at runtime ... print(f"Input: shape={x.shape}, mean={x.mean().item():.4f}") ... out = F.linear(x, linear.weight, linear.bias) ... print(f"Output: shape={out.shape}, norm={out.norm().item():.4f}") ... return (out,) ... >>> @custom_forward.register_fake ... def custom_forward_fake(linear, x): ... # Return tensor with correct shape: (batch, out_features) ... return (x.new_empty(x.shape[0], linear.weight.shape[0]),) ... >>> class MyModel(torch.nn.Module): ... def __init__(self): ... super().__init__() ... self.linear = torch.nn.Linear(10, 20) ... def forward(self, x): ... return custom_forward(self.linear, x) ... >>> model = MyModel() >>> compiled = torch.compile(model, backend="aot_eager", fullgraph=True) >>> x = torch.randn(32, 10, requires_grad=True) >>> out = compiled(x)[0] >>> out.sum().backward() # Gradients flow to model.linear.weight/bias and x Args: fn: The function being decorated. mutates_args: Set of Python expressions (as strings) identifying arguments that the function mutates in-place. Each string is evaluated against the function's parameters. Examples: ``'buf'`` for a plain tensor, ``'model.running_mean'`` for an nn.Module buffer, ``'buffers'`` to mark all tensors in a list, ``'buffers[0]'`` for a specific element, ``'state["key"]'`` for a dict entry. """ if fn is None: return functools.partial(leaf_function, mutates_args=mutates_args) from . import trace_rules _check_mutually_exclusive_decorators(fn, "leaf_function") if mutates_args: import inspect import re params = set(inspect.signature(fn).parameters) for expr in mutates_args: root = re.split(r"[.\[]", expr, maxsplit=1)[0] if root not in params: raise ValueError( f"mutates_args expression '{expr}' refers to parameter '{root}', " f"which is not a parameter of '{fn.__name__}'. " f"Available parameters: {', '.join(params)}" ) @functools.wraps(fn) def inner(*args: _P.args, **kwargs: _P.kwargs) -> _R: if inner._torchdynamo_leaf_fake_fn is None: # type: ignore[attr-defined] raise ValueError( f"leaf_function '{getattr(fn, '__name__', fn)}' " "requires a fake implementation. Please provide one using the @.register_fake " "decorator. See the leaf_function docstring for details." ) return _invoke_leaf_function_python( fn, # pyrefly: ignore [bad-argument-type] # pyrefly: ignore [bad-argument-type] inner._torchdynamo_leaf_fake_fn, args, kwargs, mutates_args=inner._torchdynamo_leaf_mutates_args, # pyrefly: ignore [missing-attribute] hook_fn=inner._torchdynamo_leaf_hook_fn, # type: ignore[attr-defined] hook_fake_fn=inner._torchdynamo_leaf_hook_fake_fn, # type: ignore[attr-defined] ) # type: ignore[attr-defined] inner._torchdynamo_leaf_real_fn = fn # type: ignore[attr-defined] inner._torchdynamo_leaf_fake_fn = None # type: ignore[attr-defined] inner._torchdynamo_leaf_mutates_args = ( # pyrefly: ignore [missing-attribute] frozenset(mutates_args) if mutates_args else frozenset() ) # type: ignore[attr-defined] inner._torchdynamo_leaf_hook_fn = None # type: ignore[attr-defined] inner._torchdynamo_leaf_hook_fake_fn = None # type: ignore[attr-defined] # Follow nonstrict_trace implementation wrapped_id = id(inner) trace_rules._allowed_callable_ids.add(wrapped_id) trace_rules._leaf_function_ids.add(wrapped_id) def deregister() -> None: trace_rules._allowed_callable_ids.remove(wrapped_id) trace_rules._leaf_function_ids.remove(wrapped_id) weakref.finalize(inner, deregister) def register_fake_setter(fake_fn: Callable[..., Any]) -> Callable[..., Any]: inner._torchdynamo_leaf_fake_fn = fake_fn # type: ignore[attr-defined] return inner inner.register_fake = register_fake_setter # type: ignore[attr-defined] def register_hook_setter(hook_fn: Callable[..., Any]) -> Callable[..., Any]: inner._torchdynamo_leaf_hook_fn = hook_fn # type: ignore[attr-defined] inner._torchdynamo_leaf_hook_fake_fn = lambda *args, **kwargs: None # type: ignore[attr-defined] return inner inner.register_multi_grad_hook = register_hook_setter # type: ignore[attr-defined] return inner def _disallow_in_graph_helper(throw_if_not_allowed: bool) -> Callable[..., Any]: def inner(fn: Any) -> Any: if isinstance(fn, (list, tuple)): return [disallow_in_graph(x) for x in fn] if not callable(fn): raise AssertionError("disallow_in_graph expects a callable") if ( throw_if_not_allowed and trace_rules.lookup_callable(fn) != variables.TorchInGraphFunctionVariable and trace_rules.lookup(fn) != variables.TorchInGraphFunctionVariable ): raise RuntimeError( "disallow_in_graph is expected to be used on an already allowed callable (like torch.* ops). " "Allowed callables means callables that TorchDynamo puts as-is in the extracted graph." ) trace_rules._allowed_callable_ids.remove(id(fn)) trace_rules._nonstrict_trace_callable_ids.remove(id(fn)) trace_rules._disallowed_callable_ids.add(id(fn)) return fn return inner def disallow_in_graph(fn: Callable[..., Any]) -> Any: """ Customize which functions TorchDynamo will exclude in the generated graph and force a graph break on. :: torch._dynamo.disallow_in_graph(torch.sub) @torch._dynamo.optimize(...) def fn(a): x = torch.add(x, 1) x = torch.sub(x, 1) x = torch.add(x, 1) return x fn(...) Will break the graph on `torch.sub`, and give two graphs each with a single `torch.add()` op. """ return _disallow_in_graph_helper(throw_if_not_allowed=True)(fn) @_disallow_in_graph_helper(throw_if_not_allowed=False) def graph_break(msg: str = "") -> None: """Force a graph break""" # NOTE: primarily used for internal debugging purposes! @_disallow_in_graph_helper(throw_if_not_allowed=False) def skip_frame(msg: str = "") -> None: """Force a skipped frame""" @_disallow_in_graph_helper(throw_if_not_allowed=False) def step_unsupported(msg: str = "") -> None: """Force a step unsupported graph break, which results in compiling the traced FX graph so far, then skipping the rest of the frame. In order to get expected behavior, there should be at least 2 ops and a part of the code not contained in any try/with blocks.""" def forbid_in_graph(fn: Any) -> Any: """ Customize which functions TorchDynamo will assert are not present while tracing. If you want a graph break on this function instead, use disallow_in_graph. TODO(voz): We now have allow_in_graph, disallow_in_graph, forbid_in_graph - some more robust documentation would not be amiss. """ if isinstance(fn, (list, tuple)): return [forbid_in_graph(x) for x in fn] if not callable(fn): raise AssertionError("forbid_in_graph applies only to callables") # pyrefly: ignore [missing-attribute] fn._dynamo_forbidden = True return fn def substitute_in_graph( original_fn: Callable[_P, _R], *, can_constant_fold_through: bool = False, skip_signature_check: bool = False, # type that is embedded in the Python interpreter is_embedded_type: bool = False, # internal use only ) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: """ Register a polyfill handler for a function, usually a C function from the C extension, to be used in place of the original function when inlining the original function in the graph. .. note:: The polyfill handler is only used when inlining the original function. It is not used when the original function is called directly. In the eager mode, the decorated function calls the performant C function rather than the polyfill handler. The polyfill handler is a function that will be called in place of the original function when inlining the original function. The polyfill handler should have the same signature and the same behavior as the original function. Args: original_fn (callable): The original function, usually a C function, to register a polyfill handler for. can_constant_fold_through (bool, optional): Whether the polyfill handler can be constant folded through. That is, if the polyfill handler is a pure function and its arguments are constant, the result of the polyfill handler can be constant folded during the compilation. Defaults to ``False``. skip_signature_check (bool, optional): Whether to skip the signature check between the original function and the polyfill handler. Defaults to ``False``. Returns: A decorator that registers the polyfill handler for the original function. Example:: >>> # xdoctest: +SKIP("conflict with the tests: duplicate polyfill handlers") >>> import operator >>> operator.indexOf([1, 2, 3, 4, 5], 3) 2 >>> torch.compile(operator.indexOf, fullgraph=True)([1, 2, 3, 4, 5], 3) Traceback (most recent call last): ... torch._dynamo.exc.Unsupported: ... >>> @torch.compiler.substitute_in_graph(operator.indexOf) ... def indexOf(a, b, /): ... for i, item in enumerate(a): ... if item is b or item == b: ... return i ... raise ValueError("sequence.index(x): x not in sequence") >>> >>> torch.compile(operator.indexOf, fullgraph=True)([1, 2, 3, 4, 5], 3) 2 """ if not is_function(original_fn) and not ( is_embedded_type and inspect.isclass(original_fn) ): raise TypeError( f"substitute_in_graph expects a function but got {type(original_fn)!r}" ) if is_embedded_type: if not inspect.isclass(original_fn): raise TypeError( f"substitute_in_graph expects a class but got {type(original_fn)!r}" ) from .variables.builder import ITERTOOLS_POLYFILLED_TYPE_IDS, ITERTOOLS_TYPE_IDS if id(original_fn) in ITERTOOLS_TYPE_IDS: ITERTOOLS_POLYFILLED_TYPE_IDS.add(id(original_fn)) def wrapper(traceable_fn: Callable[_P, _R]) -> Callable[_P, _R]: if not is_function(traceable_fn): raise TypeError( f"@substitute_in_graph(...) expects a function but got {type(traceable_fn)!r}" ) if not skip_signature_check: try: original_sig = inspect.signature(original_fn) except ValueError: pass else: traceable_sig = inspect.signature(traceable_fn) def sig_ident( sig: inspect.Signature, ) -> tuple[tuple[str, ...], set[str], dict[str, Any]]: # Ignore annotations for parameters and return type return ( tuple( p.name for p in sig.parameters.values() if ( p.kind not in { p.KEYWORD_ONLY, # the name of *args and **kwargs is not important p.VAR_POSITIONAL, p.VAR_KEYWORD, } ) ), { p.name for p in sig.parameters.values() if p.kind == p.KEYWORD_ONLY }, { p.name: p.default for p in sig.parameters.values() # the name of *args and **kwargs is not important if p.kind not in {p.VAR_POSITIONAL, p.VAR_KEYWORD} }, ) wildcard_sig = inspect.signature(lambda *args, **kwargs: None) if ( sig_ident(original_sig) != sig_ident(traceable_sig) and sig_ident(original_sig) != sig_ident(wildcard_sig) and sig_ident(traceable_sig) != sig_ident(wildcard_sig) ): raise TypeError( f"Signature mismatch between {original_fn} and {traceable_fn}: " f"{original_sig} != {traceable_sig}" ) from torch._dynamo.guards import GuardBuilder from torch._dynamo.trace_rules import ( _polyfilled_function_ids, get_torch_obj_rule_map, ) from torch._dynamo.variables import PolyfilledFunctionVariable from torch._dynamo.variables.builder import VariableBuilder id_dispatch_map = VariableBuilder._id_dispatch() if id(original_fn) in id_dispatch_map: raise ValueError( f"Duplicate dispatch rule for {original_fn}: " "already registered in VariableBuilder's id dispatch map" ) if id(original_fn) in _polyfilled_function_ids: raise ValueError(f"Duplicate polyfilled object {original_fn}") rule_map: dict[Any, type[VariableTracker]] = get_torch_obj_rule_map() if original_fn in rule_map: raise ValueError( f"Duplicate object {original_fn} with different rules: " f"{PolyfilledFunctionVariable}, {rule_map[original_fn]}" ) polyfill_handlers: dict[Callable[..., Any], FunctionType] polyfill_handlers = PolyfilledFunctionVariable._get_polyfill_handlers() if original_fn in polyfill_handlers: raise ValueError( f"Duplicate polyfill handlers for {original_fn}: " f"already handled by {polyfill_handlers[original_fn]}" ) # Need to wrap the function because we may cannot assign __torch_dynamo_polyfill__ to a # C++ function. @functools.wraps(traceable_fn) def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R: return original_fn(*args, **kwargs) def dispatch_fn( self: VariableBuilder, value: Callable[_P, _R] ) -> PolyfilledFunctionVariable: if inspect.isclass(value): guard_type = GuardBuilder.CLASS_MATCH elif inspect.ismodule(value): guard_type = GuardBuilder.MODULE_MATCH else: guard_type = GuardBuilder.ID_MATCH guards = self.install_guards(guard_type) if guards is None: raise AssertionError("install_guards returned None") return PolyfilledFunctionVariable( value, source=self.source, **guards, ) id_dispatch_map[id(original_fn)] = id_dispatch_map[id(wrapped)] = dispatch_fn _polyfilled_function_ids.add(id(original_fn)) _polyfilled_function_ids.add(id(wrapped)) rule_map[original_fn] = rule_map[wrapped] = PolyfilledFunctionVariable polyfill_handlers[original_fn] = polyfill_handlers[wrapped] = wrapped # type: ignore[assignment] wrapped.__torch_dynamo_original__ = original_fn # type: ignore[attr-defined] wrapped.__torch_dynamo_polyfill__ = traceable_fn # type: ignore[attr-defined] wrapped.__torch_dynamo_can_constant_fold_through__ = can_constant_fold_through # type: ignore[attr-defined] return wrapped # type: ignore[return-value] return wrapper # Helper function to flatten a tensor subclass and apply a function to # all inner tensors that match the outer dim. Used to reduce duplication # across the various marking APIs. def _apply_func_to_inner_tensors_of_same_dim( func: Callable[..., Any], t: object, *args: Any, **kwargs: Any ) -> None: if not is_traceable_wrapper_subclass(t): raise AssertionError(f"Expected a traceable wrapper subclass, got {type(t)}") attrs, _ctx = t.__tensor_flatten__() if not isinstance(t, torch.Tensor): raise AssertionError(f"Expected a Tensor, got {type(t)}") for attr in attrs: match getattr(t, attr): case torch.Tensor() as inner: if inner.dim() == t.dim(): func(inner, *args, **kwargs) case OpaqueBase(): pass case unexpected: raise AssertionError( f"expected Tensor or OpaqueBase, got {type(unexpected)}" ) @dataclass(frozen=True, slots=True) class _DimRange: """ This represents a dimension of a tensor and the corresponding min and max values it can take. Don't create this class directly; instead, use :func:`mark_dynamic`. """ dim: int min: int max: int @forbid_in_graph def mark_unbacked( t: Any, index: int | list[Any] | tuple[Any], hint_override: int | None = None, strict: bool = False, specialize_on: list[Any] | None = None, shape_id: str | None = None, min: int | None = None, max: int | None = None, ) -> None: """ Mark a tensor as having an unbacked dimension. This changes the semantics of operations: - The size of the specified dimension will always be reported as not equal to zero or one. - Assertions on this index will be turned into runtime asserts. - Attempting to get the real value of this dimension will raise an exception. - In effect, this dimension is treated as data-dependent (its value is unknown). Args: t (Any): The tensor to mark as having an unbacked dimension. index (int or list/tuple of int): The dimension(s) to mark as unbacked. Can be a single integer or a list/tuple of integers. Calls are additive: mark_unbacked(x, 0) followed by mark_unbacked(x, 1) marks both dims. hint_override (Optional[int], default=None): An optional integer to override the size hint for this dimension. This is only used by the inductor backend for size hint queries, such as during autotuning. NOTE: changing hint_override values will cause FxGraphCache misses, since hint overrides affect inductor codegen decisions and are included in the cache key via ShapeEnv.var_to_hint_override. strict (bool, default=False): If True, an error will be raised if the unbacked dimension is specialized. By default (strict=False), specialization is allowed and will proceed without error. specialize_on (Optional[list[Any]], default=None): A list of specialization criteria (e.g., lambdas) for this dimension. If provided, Dynamo will generate specialized compiled regions for each criterion in addition to a generic trace. shape_id (Optional[str], default=None): An optional identifier to group unbacked dimensions together. All unbacked dimensions with the same shape_id will share the same unbacked symbol. This is useful when multiple tensors are known to have the same batch size at runtime. A runtime assertion is added to ensure this property at runtime. min (Optional[int], default=None): Minimum value constraint for this dimension. If provided, a runtime check will be added to ensure the dimension is >= min. max (Optional[int], default=None): Maximum value constraint for this dimension. If provided, a runtime check will be added to ensure the dimension is <= max. """ if torch.distributed.is_available() and isinstance( t, torch.distributed.tensor.DTensor ): # apply on inner tensor sizes/strides mark_unbacked(t._local_tensor, index, shape_id=shape_id) else: # You could have copied the mark_dynamic behavior but I'm not convinced # it's what you want if is_traceable_wrapper_subclass(t): raise AssertionError("not implemented yet") if isinstance(index, int): if strict: if not hasattr(t, "_dynamo_strict_unbacked_indices"): t._dynamo_strict_unbacked_indices = set() t._dynamo_strict_unbacked_indices.add(index) t._has_dynamo_dim_marking = True # type: ignore[attr-defined] return if not hasattr(t, "_specialized_on"): # pyrefly: ignore [implicit-any] t._specialize_on = {} if not hasattr(t, "_dynamo_unbacked_indices"): t._dynamo_unbacked_indices = set() if not hasattr(t, "_dynamo_hint_overrides"): # pyrefly: ignore [implicit-any] t._dynamo_hint_overrides = {} if hint_override: t._dynamo_hint_overrides[index] = hint_override if min is not None or max is not None: if not hasattr(t, "_dynamo_unbacked_bounds"): # pyrefly: ignore [implicit-any] t._dynamo_unbacked_bounds = {} t._dynamo_unbacked_bounds[index] = (min, max) if shape_id is not None: if not hasattr(t, "_dynamo_shape_ids"): # pyrefly: ignore [implicit-any] t._dynamo_shape_ids = {} t._dynamo_shape_ids[index] = shape_id # FX tracers don't respect @forbid_in_graph and choke on the following error since it passes in proxies: # TypeError: 'Attribute' object does not support item assignment if isinstance(t._specialize_on, dict): t._specialize_on[index] = specialize_on if specialize_on is not None else [] t._dynamo_unbacked_indices.add(index) t._has_dynamo_dim_marking = True # type: ignore[attr-defined] return if not isinstance(index, (list, tuple)): raise AssertionError( f"Expected index to be int, list, or tuple, got {type(index)}" ) for i in index: mark_unbacked(t, i, shape_id=shape_id, min=min, max=max) @forbid_in_graph def mark_dynamic( t: Any, index: int | list[Any] | tuple[Any], *, hint_override: int | None = None, min: int | None = None, max: int | None = None, specialize_on: list[Any] | None = None, ) -> None: """ Mark a tensor as having a dynamic dim and set corresponding min and max range for the dim. [Note - on the state of mark_dynamic] The behavior of having a dynamic dimension on a tensor is governed by a few factors: 1) torch._dynamo.config dynamic_shapes True or False. a) dynamic_shapes=True - dynamic_shapes must be True for mark_dynamic to work. a) dynamic_shapes=False - This config will raise an exception when used in conjunction with mark_dynamic. We will eventually support this. 2) If the dimension is fully constrained - as in, it does not allow more than a single value in both eager (torch.compile, torch._dynamo.optimize) mode and export mode (torch._dynamo.export), we will raise an error 3) If the dimension is partially constrained - allowing at least 2 values but not the full unbounded range of shapes, in eager we will pass it through, but export will raise an error. 4) Attempts to trace this function will explicitly raise. As such, all calls to mark_dynamic must be made before torch.compile. 5) If hint_override is passed, the hint_override for the specified dimension will replace the provided value from the first example input as the official size hint. Note: changing hint_override values will cause FxGraphCache misses, since hint overrides affect inductor codegen decisions (autotuning, reduction strategy, etc.) and are included in the cache key via ShapeEnv.var_to_hint_override. 6) If specialize_on is passed in, we will perform a single generic Dynamo trace followed by multiple specialized compilations in addition to a single generic compilation. NB: For now we only support per dimension specialization, or in other words we do not generate a cross product of specializations. At runtime, we will dispatch to a specialized compiled region if the input matches the specialization criteria. For example: mark_dynamic(..., specialize_on=[ lambda x: x == 8, lambda x: x == 16 ]) This approach results in one Dynamo trace and two backend compilations. When the input dimension equals 8 or 16 at runtime, execution will be directed to the specialized compiled region. Performance measurements indicate 2-8x speedups depending on the specific specialization and model architecture. """ if is_traceable_wrapper_subclass(t): # default behavior: mirror mark_dynamic() on all inner tensors with same dim as t # TODO: Make this configurable via a supported public API _apply_func_to_inner_tensors_of_same_dim( mark_dynamic, t, index, min=min, max=max ) if isinstance(index, int): if not hasattr(t, "_dynamo_dynamic_indices"): t._dynamo_dynamic_indices = set() t._dynamo_dynamic_range = set() # pyrefly: ignore [implicit-any] t._dynamo_hint_overrides = {} if not hasattr(t, "_specialize_on"): # pyrefly: ignore [implicit-any] t._specialize_on = {} if hint_override: t._dynamo_hint_overrides[index] = hint_override # TODO(voz): Should we bounds check? t._dynamo_dynamic_indices.add(index) t._dynamo_dynamic_range.add(_DimRange(index, min, max)) # type: ignore[arg-type] t._has_dynamo_dim_marking = True # type: ignore[attr-defined] # FX tracers don't respect @forbid_in_graph and choke on the following error since it passes in proxies: # TypeError: 'Attribute' object does not support item assignment if isinstance(t._specialize_on, dict): t._specialize_on[index] = specialize_on if specialize_on is not None else [] return if not isinstance(index, (list, tuple)): raise AssertionError( f"Expected index to be int, list, or tuple, got {type(index)}" ) for i in index: mark_dynamic(t, i, min=min, max=max, specialize_on=specialize_on) @forbid_in_graph def maybe_mark_dynamic(t: Any, index: int | list[Any] | tuple[Any]) -> None: """ Mark a tensor as having a dynamic dim, but don't enforce it (i.e., if this dimension ends up getting specialized, don't error). """ if is_traceable_wrapper_subclass(t): # default behavior: mirror maybe_mark_dynamic() on all inner tensors with same dim as t # TODO: Make this configurable via a supported public API _apply_func_to_inner_tensors_of_same_dim(maybe_mark_dynamic, t, index) if isinstance(index, int): if not hasattr(t, "_dynamo_weak_dynamic_indices"): t._dynamo_weak_dynamic_indices = set() # TODO(voz): Should we bounds check? t._dynamo_weak_dynamic_indices.add(index) t._has_dynamo_dim_marking = True # type: ignore[attr-defined] return if not isinstance(index, (list, tuple)): raise AssertionError( f"Expected index to be int, list, or tuple, got {type(index)}" ) for i in index: maybe_mark_dynamic(t, i) def mark_static(t: Any, index: int | list[Any] | tuple[Any] | None = None) -> None: """ Mark a tensor as having a static dim or mark a nn module class as static. For tensors =========== This will prevent us from attempting to compile it dynamically when dynamic=True; this can improve trace-time performance. This has lower precedence than mark_dynamic. Unlike mark_dynamic, this can be done inside a graph, in which case it induces specialization on the tensor. For nn.Module classes ===================== For static nn.Module classes, TorchDynamo assumes that the module instance attributes will not be modified after compilation. This will ensure that TorchDynamo keeps integer attributes CONSTANT and not symints. From TorchDynamo implementation side, the instances of static-marked nn.Module class will be converted to UnspecializedBuiltinNNModuleVariable, which have the same properties. Note that we still have to guard on the attributes, because different instances of the nn.Module can have different values of the attributes. The key point here is that the attributes are static. """ if is_compiling(): if index is None: for s in t.size(): comptime.force_static(s) else: comptime.force_static(t.size(index)) return if is_traceable_wrapper_subclass(t): # default behavior: mirror mark_static() on all inner tensors with same dim as t # TODO: Make this configurable via a supported public API _apply_func_to_inner_tensors_of_same_dim(mark_static, t, index) if not isinstance(t, torch.Tensor) and issubclass(t, torch.nn.Module): t._dynamo_marked_static = True # pyrefly: ignore [bad-return] return t if not isinstance(t, torch.Tensor): raise TypeError( f"mark_static expects a tensor/nn.Module class but received {type(t)}" ) if isinstance(index, int): if not hasattr(t, "_dynamo_static_indices"): t._dynamo_static_indices = set() # type: ignore[attr-defined] # TODO(voz): Should we bounds check? t._dynamo_static_indices.add(index) # type: ignore[attr-defined] t._has_dynamo_dim_marking = True # type: ignore[attr-defined] elif index is None: for i in range(t.dim()): mark_static(t, i) else: if not isinstance(index, (list, tuple)): raise AssertionError( f"Expected index to be int, list, or tuple, got {type(index)}" ) for i in index: mark_static(t, i) @forbid_in_graph def mark_static_address(t: Any, guard: bool = False) -> None: """ Marks an input tensor whose address should be treated as constant across calls to the same dynamo-compiled function. This indicates to cudagraphs that an extra allocation is not needed for this input. The data_ptr will be guarded if guard=True, and cause a full recompile if the data_ptr changes. Note: If this address changes, cudagraphs will re-record if guard=False. """ if not isinstance(t, torch.Tensor): raise TypeError(f"mark_static_address expects a tensor but received {type(t)}") if guard: t._dynamo_static_input_type = "guarded" # type: ignore[attr-defined] else: t._dynamo_static_input_type = "unguarded" # type: ignore[attr-defined] def _patch_einops_symint_compat(einops_mod: Any) -> None: """Backport the SymInt lru_cache fix from einops 0.7.0 into einops <= 0.6.1.""" for name in ("_reconstruct_from_shape", "_prepare_transformation_recipe"): cached = getattr(einops_mod, name) uncached = cached.__wrapped__ def make_wrapper(cached_fn: Any, uncached_fn: Any) -> Any: @functools.wraps(cached_fn) def wrapper(*args: Any, **kwargs: Any) -> Any: try: return cached_fn(*args, **kwargs) except TypeError: return uncached_fn(*args, **kwargs) return wrapper setattr(einops_mod, name, make_wrapper(cached, uncached)) _EINOPS_DYNAMO_TRACING_MIN_VERSION = Version("0.8.2") # Known pure parser/shape helpers in einops 0.8.2. Keep this explicit so # unrelated or future lru_cache wrappers continue through the normal warning path. _EINOPS_LRU_CACHE_WRAPPER_ALLOWLIST = ( ( "einops.einops", ( "_reconstruct_from_shape", "_prepare_transformation_recipe", "_compactify_pattern_for_einsum", ), ), ("einops.packing", ("analyze_pattern",)), ) def _einops_supports_dynamo_tracing(einops_mod: Any) -> bool: try: return Version(einops_mod.__version__) >= _EINOPS_DYNAMO_TRACING_MIN_VERSION except InvalidVersion: return False def _allow_lru_cache_trace_without_warning_for_einops() -> None: import importlib for module_name, attr_names in _EINOPS_LRU_CACHE_WRAPPER_ALLOWLIST: try: module = importlib.import_module(module_name) except ImportError: continue for attr_name in attr_names: obj = getattr(module, attr_name, None) if is_lru_cache_wrapped_function(obj): allow_lru_cache_wrapper_trace_without_warning(obj) # Dynamo can trace through einops 0.8.2+ directly (no allow_in_graph needed). # Older versions still need the allow_in_graph registration below. def _allow_in_graph_einops() -> None: import einops if _einops_supports_dynamo_tracing(einops): if hasattr(einops, "einops") and hasattr(einops.einops, "get_backend"): # trigger backend registration up front to avoid a later guard failure # that would otherwise cause a recompilation einops.rearrange(torch.empty(1), "i -> i") # einops uses lru_cache for pure library-internal recipe helpers. Mark # only those wrappers so general lru_cache warnings are unchanged. _allow_lru_cache_trace_without_warning_for_einops() return try: # requires einops > 0.6.1, torch >= 2.0 from einops._torch_specific import ( # type: ignore[attr-defined] # noqa: F401 _ops_were_registered_in_torchdynamo, ) # einops > 0.6.1 will call the op registration logic as it is imported. except ImportError: # einops <= 0.6.1 doesn't handle unhashable SymInt in its lru_cache'd # helpers. Backport the try/except TypeError fallback from einops 0.7.0+ # so allow_in_graph works during fake tensor validation. _patch_einops_symint_compat(einops.einops) # type: ignore[attr-defined] allow_in_graph(einops.rearrange) allow_in_graph(einops.reduce) if hasattr(einops, "repeat"): allow_in_graph(einops.repeat) # available since einops 0.2.0 if hasattr(einops, "einsum"): allow_in_graph(einops.einsum) # available since einops 0.5.0 if hasattr(einops, "pack"): allow_in_graph(einops.pack) # available since einops 0.6.0 if hasattr(einops, "unpack"): allow_in_graph(einops.unpack) # available since einops 0.6.0 # Note: this carefully avoids eagerly import einops. trace_rules.add_module_init_func("einops", _allow_in_graph_einops) # Proxy class for torch._dynamo.config patching - so dynamo can identify context managers/decorators # created by patch_dynamo_config, compared to ones created by a raw torch._dynamo.config.patch. class DynamoConfigPatchProxy: def __init__(self, config_patch: Any) -> None: self.config_patch = config_patch @property def changes(self) -> dict[str, Any]: return self.config_patch.changes # Decorator implementation that simply sets up `self` as a context manager. # Placed in external_utils so that we can trace through it. __call__ = wrap_dunder_call_ctx_manager def __enter__(self) -> None: return self.config_patch.__enter__() def __exit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None, ) -> None: return self.config_patch.__exit__(exc_type, exc_val, exc_tb) # Criteria for patchable config: # - Config values must be constants (i.e. int, float, str, bool, None). # - in particular, NO list, set, dict. # - Traceable config patches are only useful for configs that change dynamo behavior # from symbolic_convert and below. # - e.g. patching recompile_limit won't really do anything. # - For patching configs that affect Dynamo behavior above symbolic_convert, # ensure that Dynamo behaves soundly even if tracing is done with different config. # - e.g. be careful if patching guard-related configs as configs may have changed # between guard creation and evaluation. _allowed_config_patches = ( "verbose", "verify_correctness", "rewrite_assert_with_torch_assert", "capture_scalar_outputs", "allow_unspec_int_on_nn_module", "skip_torchrec", "dont_skip_tracing", "nested_graph_breaks", ) from . import config for name in _allowed_config_patches: if not hasattr(config, name): raise AssertionError(f"nonexistent config: {name}") del config def _patch_dynamo_config_check(changes: dict[str, Any]) -> None: for k, v in changes.items(): if k not in _allowed_config_patches: raise ValueError( f"patch_dynamo_config does not support patching config {k}" ) if not torch._dynamo.utils.is_safe_constant(v): raise ValueError( f"patch_dynamo_config does not support patching config {k} " f"with non-safe-constant value {v}" ) # TODO: also implement nonrecursive patch_dynamo_config/dont_skip_tracing. # Unlike config.patch, we also need to accept tuple as input in order to # deal with context manager reconstruction. def patch_dynamo_config( arg1: str | dict[str, Any] | tuple[tuple[str, Any], ...] | None = None, arg2: Any = None, **kwargs: Any, ) -> DynamoConfigPatchProxy: """ A wrapper around torch._dynamo.config.patch that can be traced by Dynamo to temporarily change config values DURING tracing. See _allowed_config_patches for the list of allowed config patches. Arguments are the same as with torch._dynamo.config.patch. Can be used as a decorator or a context manager. User code SHOULD NOT MODIFY the return value of this function. WARNING: changing Dynamo config during tracing can lead to unpredictable tracing behavior! Proceed only as advised! """ if isinstance(arg1, tuple): arg1 = dict(arg1) config_patch = torch._dynamo.config.patch(arg1, arg2, **kwargs) _patch_dynamo_config_check(config_patch.changes) # check for valid patching using config_patch.changes return DynamoConfigPatchProxy(config_patch) @overload def dont_skip_tracing(fn: None = None) -> DynamoConfigPatchProxy: ... @overload def dont_skip_tracing(fn: Callable[_P, _R]) -> Callable[_P, _R]: ... def dont_skip_tracing(fn: Any | None = None) -> Any: """ Context manager/decorator to trace into functions intentionally marked by developers to be skipped when tracing. This decorator will also apply to recursively invoked functions. """ ctx = patch_dynamo_config(dont_skip_tracing=True) if fn: return ctx(fn) return ctx @overload def disable_nested_graph_breaks(fn: None = None) -> DynamoConfigPatchProxy: ... @overload def disable_nested_graph_breaks(fn: Callable[_P, _R]) -> Callable[_P, _R]: ... def disable_nested_graph_breaks(fn: Any | None = None) -> Any: """ Context manager/decorator to disable nested graph breaks when tracing this function and any nested functions. Used when nested graph breaks is causing problems. """ ctx = patch_dynamo_config(nested_graph_breaks=False) if fn: return ctx(fn) return ctx class ErrorOnGraphBreakDecoratorContextManager: def __init__(self, error_on_graph_break: bool) -> None: self.error_on_graph_break = error_on_graph_break self.prev_error_on_graph_break: list[bool | None] = [] __call__ = wrap_dunder_call_ctx_manager def __enter__(self) -> None: current_error_on_graph_break = _get_error_on_graph_break() if current_error_on_graph_break != self.error_on_graph_break: self.prev_error_on_graph_break.append(current_error_on_graph_break) _set_error_on_graph_break(self.error_on_graph_break) else: self.prev_error_on_graph_break.append(None) def __exit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None, ) -> None: prev_error_on_graph_break = self.prev_error_on_graph_break.pop() if prev_error_on_graph_break is not None: _set_error_on_graph_break(prev_error_on_graph_break) def error_on_graph_break( error_on_graph_break: bool, ) -> ErrorOnGraphBreakDecoratorContextManager: """ Context manager/decorator to toggle torch.compile's `error_on_graph_break` setting at compile time. If `fullgraph` is set, then `error_on_graph_break` does nothing (i.e. `fullgraph = True` takes higher precedence). If `fullgraph` is False, then `error_on_graph_break` determines whether `torch.compile` throws an error upon encountering a graph break, or attempts to continue tracing. `error_on_graph_break` can be toggled during compile time with this decorator to allow graph breaks in some compiled regions but not others. One key difference from `fullgraph` is that `error_on_graph_break = True` does NOT guarantee that a single graph is captured from the compiled function. The default value of torch.compile's `error_on_graph_break` setting is False. """ return ErrorOnGraphBreakDecoratorContextManager(error_on_graph_break) class CudagraphOverrideContextManager: """Context manager that overrides cudagraph recording during tracing.""" def __init__(self, fwd: bool | None = None, bwd: bool | None = None) -> None: self.fwd = fwd self.bwd = bwd __call__ = wrap_dunder_call_ctx_manager def __enter__(self) -> None: pass def __exit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None, ) -> None: pass def override_cudagraphs( fwd: bool | None = None, bwd: bool | None = None ) -> CudagraphOverrideContextManager: """ Context manager/decorator to override cudagraph recording for compiled graphs. When used as a context manager, overrides cudagraphs for all graph segments within the block (including across graph breaks). When used as a decorator, marks a function so that any compiled graph inlining it will have cudagraphs overridden. Args: fwd: If False, disable cudagraphs for forward. If True, force enable. If None, don't override. bwd: If False, disable cudagraphs for backward. If True, force enable. If None, don't override. """ return CudagraphOverrideContextManager(fwd=fwd, bwd=bwd) def override_optimization_hint(x: Any, val: int) -> None: """Override the optimization hint for a scalar unbacked symbol. When the compiler or runtime needs a non-guarding integer hint for an unbacked ``SymInt`` — for example during FX passes, graph partitioning, or inductor autotuning — it calls ``_optimization_hint_base`` (see ``torch/fx/experimental/_size_hinting.py``). By default that function uses internal heuristics to choose a hint and a global fixed fallback; this function lets user code override that choice. This is similar to the ``hint_override`` parameter in ``mark_unbacked``, but applies to symbols that already exist (e.g. from ``.item()`` calls). Typical usage:: u = x.item() # unbacked SymInt torch._dynamo.override_optimization_hint(u, 42) # From now on, any call to shape_env.optimization_hint(u, ...) # returns 42 instead of the default heuristic value. This updates ``shape_env.var_to_hint_override`` so that any consumer of ``_optimization_hint_base`` sees *val* as the hint for the unbacked symbol behind *x*. Works both eagerly (during FX passes or outside dynamo) and inside ``torch.compile`` regions. Behavior during compilation: The dynamo handler applies the hint as a **side effect** on ``shape_env.var_to_hint_override`` during tracing. No FX graph node is emitted — the call is fully consumed at trace time and does not appear in the pre-grad, joint, or post-autograd graphs. ``FXGraphCache`` includes ``var_to_hint_override`` in its cache key, so cache hits/misses correctly reflect hint changes. .. note:: To maximize performance, it is recommended to pass hints for **all** unbacked symbols in the program to guide optimizations. Args: x: A ``torch.SymInt`` wrapping an unbacked symbol (e.g. from ``.item()``), or a plain ``int``. If *x* is a plain ``int`` the call is a no-op. val: The integer hint value to record. """ if not isinstance(val, int): raise TypeError( f"override_optimization_hint expects val to be an int, got {type(val)}" ) if isinstance(x, int): return if not isinstance(x, torch.SymInt): raise TypeError( f"override_optimization_hint expects a torch.SymInt or int, got {type(x)}" ) shape_env = x.node.shape_env expr = x.node.expr import sympy if not isinstance(expr, sympy.Symbol): raise ValueError( f"override_optimization_hint expects a single unbacked symbol, " f"got derived expression: {expr}" ) if not shape_env.is_unbacked_symint(expr): raise ValueError( f"override_optimization_hint expects an unbacked symbol, " f"but {expr} is backed" ) shape_env.var_to_hint_override[expr] = val def is_dynamo_disable_recursive(method: Callable[[Any], Any]) -> bool | None: """ Check if a method is marked as `dynamo_disable` recursively. It returns: - True if disable(recursive=True) - False if disable(recursive=False) - None if method is not a disable decorator """ return getattr(method, "_torchdynamo_disable_recursive", None) _HASH_SLOTS = ("__hash__",) _RICHCOMPARE_SLOTS = ("__eq__", "__ne__", "__lt__", "__le__", "__gt__", "__ge__") def allow_c_slot( tp: type, *, tp_hash: bool = True, tp_richcompare: bool = True, ) -> type: """Register a C extension type's slots as safe to call at trace time. By default, ``torch.compile`` graph-breaks when it encounters ``hash()`` or comparison operators on a C extension type with custom C slots (e.g., types defined in C extension modules). This function tells Dynamo that the type's C slots are safe to evaluate during tracing, avoiding graph breaks. The slot functions must satisfy these requirements: - **Pure**: depends only on the object's value, with no observable side effects (no I/O, no mutation of global state). - **Deterministic**: returns the same result for the same object across calls within a process. - **Immutable objects**: the object's state must not change after construction in a way that affects the slot's return value. Builtin types (``int``, ``str``, etc.) and Python-level dunder methods are already handled and do not need registration. This API is only needed for C extension types whose C slots Dynamo cannot trace into. Args: tp: The C extension type whose C slots should be allowed. tp_hash: Register ``__hash__`` as safe (default True). tp_richcompare: Register comparison dunders as safe (default True). Returns: The type, unchanged (so it can be used as a decorator). Example:: import torch._dynamo from my_extension import MyType # Register all slots (hash + comparison) torch._dynamo.allow_c_slot(MyType) # Register only hash torch._dynamo.allow_c_slot(MyType, tp_richcompare=False) # Register only comparison torch._dynamo.allow_c_slot(MyType, tp_hash=False) """ from .variables.user_defined import _safe_c_slots if not isinstance(tp, type): raise TypeError(f"allow_c_slot expects a type, got {type(tp).__name__}") safe = _safe_c_slots() dunders = () if tp_hash: dunders += _HASH_SLOTS if tp_richcompare: dunders += _RICHCOMPARE_SLOTS for dunder in dunders: fn = getattr(tp, dunder, None) if fn is not None and fn is not getattr(object, dunder, None): safe.add(fn) return tp