| import functools |
| import inspect |
| import itertools |
| import logging |
| import math |
| import operator |
| import types |
| from typing import Dict, List |
| |
| import torch |
| from torch import sym_float, sym_int |
| |
| from .. import config, variables |
| from ..allowed_functions import is_allowed |
| from ..exc import ( |
| AttributeMutationError, |
| unimplemented, |
| Unsupported, |
| UserError, |
| UserErrorType, |
| ) |
| from ..guards import GuardBuilder |
| from ..replay_record import DummyModule |
| from ..source import AttrSource, is_constant_source, SuperSource, TypeSource |
| from ..utils import ( |
| build_checkpoint_variable, |
| check_constant_args, |
| check_numpy_ndarray_args, |
| check_unspec_python_args, |
| get_fake_value, |
| guard_if_dyn, |
| is_utils_checkpoint, |
| istype, |
| numpy_operator_wrapper, |
| proxy_args_kwargs, |
| specialize_args_kwargs, |
| ) |
| from .base import MutableLocal, typestr, VariableTracker |
| from .constant import ConstantVariable, EnumVariable |
| from .dicts import ConstDictVariable |
| from .lists import ( |
| BaseListVariable, |
| ListIteratorVariable, |
| ListVariable, |
| SetVariable, |
| SizeVariable, |
| TupleIteratorVariable, |
| TupleVariable, |
| ) |
| from .tensor import FakeItemVariable, SymNodeVariable, UnspecializedPythonVariable |
| from .user_defined import UserDefinedVariable |
| |
| log = logging.getLogger(__name__) |
| |
| |
| class BuiltinVariable(VariableTracker): |
| @staticmethod |
| @functools.lru_cache(None) |
| def _constant_fold_functions(): |
| fns = { |
| abs, |
| all, |
| any, |
| bool, |
| callable, |
| chr, |
| divmod, |
| float, |
| int, |
| len, |
| max, |
| min, |
| ord, |
| pow, |
| repr, |
| round, |
| str, |
| str.format, |
| sum, |
| type, |
| operator.pos, |
| operator.neg, |
| operator.not_, |
| operator.invert, |
| operator.pow, |
| operator.mul, |
| operator.matmul, |
| operator.floordiv, |
| operator.truediv, |
| operator.mod, |
| operator.add, |
| operator.sub, |
| operator.getitem, |
| operator.lshift, |
| operator.rshift, |
| operator.and_, |
| operator.or_, |
| operator.xor, |
| operator.ipow, |
| operator.imul, |
| operator.imatmul, |
| operator.ifloordiv, |
| operator.itruediv, |
| operator.imod, |
| operator.iadd, |
| operator.isub, |
| operator.ilshift, |
| operator.irshift, |
| operator.iand, |
| operator.ixor, |
| operator.ior, |
| operator.index, |
| } |
| fns.update(x for x in math.__dict__.values() if isinstance(x, type(math.sqrt))) |
| return fns |
| |
| def can_constant_fold_through(self): |
| return self.fn in self._constant_fold_functions() |
| |
| @staticmethod |
| @functools.lru_cache(None) |
| def _fx_graph_functions(): |
| fns = { |
| operator.pos, |
| operator.neg, |
| operator.not_, |
| operator.invert, |
| operator.pow, |
| operator.mul, |
| operator.matmul, |
| operator.floordiv, |
| operator.truediv, |
| operator.mod, |
| operator.add, |
| operator.sub, |
| operator.getitem, |
| operator.lshift, |
| operator.rshift, |
| operator.and_, |
| operator.or_, |
| operator.xor, |
| operator.ipow, |
| operator.imul, |
| operator.imatmul, |
| operator.ifloordiv, |
| operator.itruediv, |
| operator.imod, |
| operator.iadd, |
| operator.isub, |
| operator.ilshift, |
| operator.irshift, |
| operator.iand, |
| operator.ixor, |
| operator.ior, |
| } |
| return fns |
| |
| @staticmethod |
| @functools.lru_cache(None) |
| def _binops(): |
| # function -> ([forward name, reverse name, in-place name], in-place op) |
| fns = { |
| operator.add: (["__add__", "__radd__", "__iadd__"], operator.iadd), |
| operator.sub: (["__sub__", "__rsub__", "__isub__"], operator.isub), |
| operator.mul: (["__mul__", "__rmul__", "__imul__"], operator.imul), |
| operator.truediv: ( |
| ["__truediv__", "__rtruediv__", "__itruediv__"], |
| operator.itruediv, |
| ), |
| operator.floordiv: ( |
| ["__floordiv__", "__rfloordiv__", "__ifloordiv__"], |
| operator.ifloordiv, |
| ), |
| operator.mod: (["__mod__", "__rmod__", "__imod__"], operator.imod), |
| pow: (["__pow__", "__rpow__", "__ipow__"], operator.ipow), |
| operator.pow: (["__pow__", "__rpow__", "__ipow__"], operator.ipow), |
| operator.lshift: ( |
| ["__lshift__", "__rlshift__", "__ilshift__"], |
| operator.ilshift, |
| ), |
| operator.rshift: ( |
| ["__rshift__", "__rrshift__", "__irshift__"], |
| operator.irshift, |
| ), |
| # NB: The follow binary operators are not supported for now, since the |
| # corresponding magic methods aren't defined on SymInt / SymFloat: |
| # operator.matmul |
| # divmod |
| # operator.and_ |
| # operator.or_ |
| # operator.xor |
| } |
| return fns |
| |
| @staticmethod |
| @functools.lru_cache(None) |
| def _binop_handlers(): |
| # Multiple dispatch mechanism defining custom binop behavior for certain type |
| # combinations. Handlers are attempted in order, and will be used if the type checks |
| # match. They are expected to have the signature: |
| # fn(tx, arg0: VariableTracker, arg1: VariableTracker, options) -> VariableTracker |
| |
| # Override table contains: op_fn -> [list of handlers] |
| op_handlers = {} |
| for ( |
| op, |
| (magic_method_names, in_place_op), |
| ) in BuiltinVariable._binops().items(): |
| op_handlers[op] = [] |
| op_handlers[in_place_op] = [] |
| |
| forward_name, reverse_name, inplace_name = magic_method_names |
| |
| # User-defined args (highest precedence) |
| def user_defined_handler( |
| tx, |
| a, |
| b, |
| options, |
| forward_name=forward_name, |
| reverse_name=reverse_name, |
| ): |
| # Manually handle reversing logic if needed (e.g. call __radd__) |
| |
| # TODO: If we expand this to handle tensor args, we need to manually |
| # handle cases like this: |
| # |
| # class A(int): |
| # def __radd__(self, other): |
| # print("woof") |
| # torch.randn(3) + A(3) |
| # |
| # In this example, A.__radd__() is not called -> nothing is printed, because |
| # Tensor.__add__ only does a subtype test against int, ignoring the subclass. |
| # To be fully correct, we should not call A.__radd__() here, and there may be |
| # other cases to reason about and add exceptions for. |
| if isinstance(a, UserDefinedVariable): |
| return a.call_method(tx, forward_name, [b], {}) |
| else: |
| return b.call_method(tx, reverse_name, [a], {}) |
| |
| op_handlers[op].append( |
| ((UserDefinedVariable, VariableTracker), user_defined_handler) |
| ) |
| op_handlers[op].append( |
| ((VariableTracker, UserDefinedVariable), user_defined_handler) |
| ) |
| |
| def user_defined_inplace_handler( |
| tx, a, b, options, forward_name=inplace_name |
| ): |
| return a.call_method(tx, forward_name, [b], {}) |
| |
| op_handlers[in_place_op].append( |
| ((UserDefinedVariable, VariableTracker), user_defined_inplace_handler) |
| ) |
| op_handlers[in_place_op].append( |
| ((VariableTracker, UserDefinedVariable), user_defined_inplace_handler) |
| ) |
| |
| # Dynamic shape args |
| def dynamic_handler(tx, a, b, options, fn=op): |
| from .builder import wrap_fx_proxy |
| |
| return wrap_fx_proxy( |
| tx, |
| tx.output.create_proxy( |
| "call_function", fn, *proxy_args_kwargs([a, b], {}) |
| ), |
| **options, |
| ) |
| |
| op_handlers[op].append( |
| ((SymNodeVariable, VariableTracker), dynamic_handler) |
| ) |
| op_handlers[op].append( |
| ((VariableTracker, SymNodeVariable), dynamic_handler) |
| ) |
| |
| # NB: Prefer out-of-place op when calling in-place op to generate valid graph |
| op_handlers[in_place_op].append( |
| ((SymNodeVariable, VariableTracker), dynamic_handler) |
| ) |
| op_handlers[in_place_op].append( |
| ((VariableTracker, SymNodeVariable), dynamic_handler) |
| ) |
| |
| # Special cases - lower precedence but still prefer these over constant folding |
| |
| # List-like addition (e.g. [1, 2] + [3, 4]) |
| def tuple_add_handler(tx, a, b, options): |
| return TupleVariable(a.items + list(b.unpack_var_sequence(tx)), **options) |
| |
| def size_add_handler(tx, a, b, options): |
| return SizeVariable(a.items + list(b.unpack_var_sequence(tx)), **options) |
| |
| list_like_addition_handlers = [ |
| # NB: Prefer the tuple-specific logic over base logic because of |
| # some SizeVariable weirdness. Specifically, the tuple-specific logic |
| # drops the subclass type (e.g. SizeVariable) and returns TupleVariables. |
| ( |
| (SizeVariable, SizeVariable), |
| size_add_handler, |
| ), |
| ( |
| (TupleVariable, TupleVariable), |
| tuple_add_handler, |
| ), |
| ( |
| (TupleVariable, ConstantVariable), |
| tuple_add_handler, |
| ), |
| ( |
| (ConstantVariable, TupleVariable), |
| lambda tx, a, b, options: TupleVariable( |
| list(a.unpack_var_sequence(tx)) + b.items, **options |
| ), |
| ), |
| ( |
| (BaseListVariable, BaseListVariable), |
| lambda tx, a, b, options: type(a)(a.items + b.items, **options), |
| ), |
| ] |
| op_handlers[operator.add].extend(list_like_addition_handlers) |
| |
| def list_iadd_handler(tx, a, b, options): |
| if not a.mutable_local or not b.has_unpack_var_sequence(tx): |
| # Handler doesn't apply |
| return None |
| |
| return tx.replace_all( |
| a, |
| ListVariable( |
| list(a.items) + list(b.unpack_var_sequence(tx)), |
| regen_guards=False, |
| **options, |
| ), |
| ) |
| |
| list_like_iadd_handlers = [ |
| ( |
| (ListVariable, VariableTracker), |
| list_iadd_handler, |
| ), |
| ( |
| (TupleVariable, TupleVariable), |
| tuple_add_handler, |
| ), |
| ( |
| (TupleVariable, ConstantVariable), |
| tuple_add_handler, |
| ), |
| ] |
| op_handlers[operator.iadd].extend(list_like_iadd_handlers) |
| |
| # List-like expansion (e.g. [1, 2, 3] * 3) |
| def expand_list_like(tx, lst, const, options): |
| return lst.__class__( |
| items=lst.items * const.as_python_constant(), |
| mutable_local=MutableLocal(), |
| **options, |
| ) |
| |
| list_like_expansion_handlers = [ |
| ((ListVariable, ConstantVariable), expand_list_like), |
| ((TupleVariable, ConstantVariable), expand_list_like), |
| ( |
| (ConstantVariable, ListVariable), |
| lambda tx, a, b, options: expand_list_like(tx, b, a, options), |
| ), |
| ( |
| (ConstantVariable, TupleVariable), |
| lambda tx, a, b, options: expand_list_like(tx, b, a, options), |
| ), |
| ] |
| op_handlers[operator.mul].extend(list_like_expansion_handlers) |
| |
| return op_handlers |
| |
| @staticmethod |
| def _find_binop_handler(op, a, b): |
| handlers = BuiltinVariable._binop_handlers() |
| if op not in handlers: |
| return None |
| |
| # Return first handler that matches the type checks |
| for (type1, type2), handler in handlers[op]: |
| if isinstance(a, type1) and isinstance(b, type2): |
| return handler |
| |
| return None |
| |
| def can_insert_in_graph(self): |
| return self.fn in self._fx_graph_functions() |
| |
| def __init__(self, fn, **kwargs): |
| super().__init__(**kwargs) |
| self.fn = fn |
| |
| def __str__(self): |
| if self.fn is None: |
| name = "None" |
| else: |
| name = self.fn.__name__ |
| |
| return f"{self.__class__.__name__}({name})" |
| |
| def python_type(self): |
| return type(self.fn) |
| |
| def as_python_constant(self): |
| return self.fn |
| |
| def as_proxy(self): |
| DTYPE = { |
| bool: torch.bool, |
| int: torch.int64, |
| float: torch.float64, |
| } |
| if self.fn in DTYPE: |
| return DTYPE[self.fn] |
| return super().as_proxy() |
| |
| def reconstruct(self, codegen): |
| name = self.fn.__name__ |
| assert self.fn.__module__ == "builtins" |
| assert name not in codegen.tx.f_globals, "shadowed global" |
| return [codegen.create_load_global(name, False, add=True)] |
| |
| def constant_args(self, *args, **kwargs): |
| return check_constant_args(args, kwargs) |
| |
| def tensor_args(self, *args, **kwargs): |
| return any( |
| isinstance(i, variables.TensorVariable) |
| for i in itertools.chain(args, kwargs.values()) |
| ) and not any( |
| isinstance(i, variables.GetAttrVariable) |
| for i in itertools.chain(args, kwargs.values()) |
| ) |
| |
| def unspec_python_args(self, *args, **kwargs): |
| return check_unspec_python_args(args, kwargs) |
| |
| @staticmethod |
| def unwrap_unspec_args_kwargs(args, kwargs): |
| unwrapped_args = [] |
| unwrapped_kwargs = {} |
| for x in args: |
| if isinstance( |
| x, |
| (variables.UnspecializedPythonVariable,), |
| ): |
| unwrapped_args.append(x.raw_value) |
| else: |
| unwrapped_args.append(x.as_python_constant()) |
| for k, v in kwargs: |
| if isinstance( |
| x, |
| (variables.UnspecializedPythonVariable,), |
| ): |
| unwrapped_kwargs.update({k: v.raw_value}) |
| else: |
| unwrapped_kwargs.update({k: v.as_python_constant()}) |
| return unwrapped_args, unwrapped_kwargs |
| |
| def call_function( |
| self, tx, args: "List[VariableTracker]", kwargs: "Dict[str, VariableTracker]" |
| ) -> "VariableTracker": |
| from .builder import wrap_fx_proxy, wrap_fx_proxy_cls |
| |
| constant_args = check_constant_args(args, kwargs) |
| tensor_args = self.tensor_args(*args, **kwargs) |
| unspec_python_args = self.unspec_python_args(*args, **kwargs) |
| options = VariableTracker.propagate(self, args, kwargs.values()) |
| has_constant_handler = self.can_constant_fold_through() and ( |
| constant_args or unspec_python_args |
| ) |
| assert isinstance(args, (list, tuple)) |
| assert isinstance(kwargs, dict) |
| |
| # args[0] is list and args[1] is unspec |
| if self.fn is operator.getitem and not isinstance( |
| args[0], variables.TensorVariable |
| ): |
| tensor_args = False |
| args, kwargs = specialize_args_kwargs(tx, args, kwargs) |
| |
| if ( |
| self.can_insert_in_graph() |
| and tensor_args |
| and not ( |
| self.fn is operator.getitem |
| and isinstance(args[0], ConstDictVariable) |
| and isinstance(args[1], variables.TensorVariable) |
| ) |
| ): |
| try: |
| fn = self.fn |
| if self.fn is operator.iadd and isinstance( |
| args[0], variables.ConstantVariable |
| ): |
| # Work around weird bug in hf_T5 |
| fn, args = operator.add, [args[1], args[0]] |
| |
| if self.fn is operator.getitem and isinstance(args[1], SymNodeVariable): |
| # Standard indexing will force specialization due to |
| # __index__. Rewrite as a regular torch op which will |
| # trace fine |
| fn, args = torch.select, [ |
| args[0], |
| variables.ConstantVariable(0), |
| args[1], |
| ] |
| |
| # Interaction between ndarray and tensors: |
| # We prefer the tensor op whenever there are tensors involved |
| if check_numpy_ndarray_args(args, kwargs) and not any( |
| type(arg) == variables.TensorVariable for arg in args |
| ): |
| proxy = tx.output.create_proxy( |
| "call_function", |
| numpy_operator_wrapper(self.fn), |
| *proxy_args_kwargs(args, kwargs), |
| ) |
| |
| return wrap_fx_proxy_cls( |
| variables.NumpyNdarrayVariable, tx, proxy, **options |
| ) |
| |
| proxy = tx.output.create_proxy( |
| "call_function", |
| fn, |
| *proxy_args_kwargs(args, kwargs), |
| ) |
| if any(isinstance(arg, FakeItemVariable) for arg in args): |
| return wrap_fx_proxy_cls( |
| FakeItemVariable, |
| tx, |
| proxy, |
| **options, |
| ) |
| elif self.unspec_python_args(*args, **kwargs): |
| _args, _kwargs = self.unwrap_unspec_args_kwargs(args, kwargs) |
| raw_value = self.fn(*_args, **_kwargs) |
| |
| need_unwrap = any( |
| x.need_unwrap |
| for x in itertools.chain(args, kwargs.values()) |
| if isinstance(x, variables.UnspecializedPythonVariable) |
| ) |
| |
| return wrap_fx_proxy_cls( |
| UnspecializedPythonVariable, |
| tx, |
| proxy, |
| raw_value=raw_value, |
| need_unwrap=need_unwrap, |
| **options, |
| ) |
| elif all(isinstance(x, SymNodeVariable) for x in args): |
| return SymNodeVariable.create(tx, proxy, None, **options) |
| else: |
| # Work around for vision_maskrcnn due to precision difference |
| # specialize the dividend when float divide by tensor |
| if self.fn is operator.truediv and isinstance( |
| args[0], variables.UnspecializedPythonVariable |
| ): |
| args[0] = args[0].convert_to_constant(tx) |
| return wrap_fx_proxy(tx, proxy, **options) |
| |
| except NotImplementedError: |
| unimplemented(f"partial tensor op: {self} {args} {kwargs}") |
| |
| # Handle cases like int(torch.seed()) |
| # Also handle sym_float to sym_int cases |
| if self.fn in (int, float) and isinstance(args[0], SymNodeVariable): |
| fn_ = sym_int if self.fn is int else sym_float |
| out = wrap_fx_proxy( |
| tx=tx, |
| proxy=tx.output.create_proxy( |
| "call_function", |
| fn_, |
| (args[0].as_proxy(),), |
| {}, |
| ), |
| **options, |
| ) |
| return out |
| |
| # Handle binary ops (e.g. __add__ / __radd__, __iadd__, etc.) |
| # NB: Tensor args are handled above and not here |
| if len(kwargs) == 0 and len(args) == 2: |
| # Try to find a handler for the arg types; otherwise, fall through to constant handler |
| binop_handler = BuiltinVariable._find_binop_handler( |
| self.fn, args[0], args[1] |
| ) |
| if binop_handler: |
| res = binop_handler(tx, args[0], args[1], options) |
| if res is not None: |
| return res |
| |
| handler = getattr(self, f"call_{self.fn.__name__}", None) |
| if handler: |
| try: |
| inspect.signature(handler).bind(tx, *args, **kwargs) |
| except TypeError as exc: |
| if not has_constant_handler: |
| log.warning( |
| "incorrect arg count %s %s and no constant handler", |
| handler, |
| exc, |
| ) |
| handler = None |
| |
| if handler: |
| try: |
| result = handler(tx, *args, **kwargs) |
| if result is not None: |
| return result.add_options(options) |
| except Unsupported as exc: |
| if not has_constant_handler: |
| raise |
| # Actually, we will handle this just fine |
| exc.remove_from_stats() |
| |
| if has_constant_handler: |
| args, kwargs = specialize_args_kwargs(tx, args, kwargs) |
| # constant fold |
| return variables.ConstantVariable( |
| self.as_python_constant()( |
| *[x.as_python_constant() for x in args], |
| **{k: v.as_python_constant() for k, v in kwargs.items()}, |
| ), |
| **options, |
| ) |
| |
| if self.fn is round: |
| if len(args) > 0 and isinstance(args[0], SymNodeVariable): |
| raise UserError( |
| UserErrorType.STANDARD_LIBRARY, |
| "Calling round() on symbolic value is not supported. " |
| "You can use floor() to implement this functionality", |
| ) |
| return super().call_function(tx, args, kwargs) |
| |
| def _call_min_max(self, tx, *args): |
| if len(args) == 1 and args[0].has_unpack_var_sequence(tx): |
| # expand iterable |
| items = args[0].unpack_var_sequence(tx) |
| return self._call_min_max_seq(tx, items) |
| elif len(args) == 2: |
| return self._call_min_max_binary(tx, args[0], args[1]) |
| elif len(args) > 2: |
| return self._call_min_max_seq(tx, args) |
| |
| def _call_min_max_seq(self, tx, items): |
| assert len(items) > 0 |
| if len(items) == 1: |
| return items[0] |
| |
| return functools.reduce(functools.partial(self._call_min_max_binary, tx), items) |
| |
| def _call_min_max_binary(self, tx, a, b): |
| if self.tensor_args(a, b): |
| if not isinstance(a, variables.TensorVariable): |
| a, b = b, a |
| assert isinstance(a, variables.TensorVariable) |
| |
| # result of an item call is a scalar convert to a tensor |
| if isinstance(a, FakeItemVariable): |
| a = variables.TorchVariable(torch.tensor).call_function(tx, [a], {}) |
| |
| # Dynamic input does not get resolved, rather, gets stored as call_function |
| if isinstance(a, SymNodeVariable) or isinstance(b, SymNodeVariable): |
| from .builder import wrap_fx_proxy_cls |
| |
| return wrap_fx_proxy_cls( |
| type(a), |
| tx=tx, |
| proxy=tx.output.create_proxy( |
| "call_function", |
| self.fn, |
| *proxy_args_kwargs([a, b], {}), |
| ), |
| **VariableTracker.propagate(self, [a, b]), |
| ) |
| |
| # convert min/max to torch ops |
| if b.is_python_constant(): |
| if isinstance(a, variables.NumpyNdarrayVariable): |
| import numpy as np |
| |
| fn = variables.NumpyVariable(np.clip) |
| else: |
| fn = variables.TorchVariable(torch.clamp) |
| kwargs = {"min": b} if (self.fn is max) else {"max": b} |
| result = fn.call_function(tx, [a], kwargs) |
| else: |
| if isinstance(a, variables.NumpyNdarrayVariable): |
| import numpy as np |
| |
| fn = {max: np.maximum, min: np.minimum}[self.fn] |
| fn = variables.NumpyVariable(fn) |
| else: |
| fn = {max: torch.maximum, min: torch.minimum}[self.fn] |
| fn = variables.TorchVariable(fn) |
| result = fn.call_function(tx, [a, b], {}) |
| |
| # return unspec if both a, b are unspec or const |
| if all( |
| isinstance( |
| i, |
| ( |
| variables.UnspecializedPythonVariable, |
| variables.ConstantVariable, |
| ), |
| ) |
| for i in [a, b] |
| ): |
| if any(isinstance(val, FakeItemVariable) for val in [a, b]): |
| return variables.FakeItemVariable.from_tensor_variable(result) |
| |
| if b.is_python_constant(): |
| raw_b = b.as_python_constant() |
| else: |
| raw_b = b.raw_value |
| if self.fn is max: |
| raw_res = max(a.raw_value, raw_b) |
| else: |
| raw_res = min(a.raw_value, raw_b) |
| |
| need_unwrap = any( |
| x.need_unwrap |
| for x in [a, b] |
| if isinstance(x, variables.UnspecializedPythonVariable) |
| ) |
| return variables.UnspecializedPythonVariable.from_tensor_variable( |
| result, raw_res, need_unwrap |
| ) |
| # otherwise return tensor |
| else: |
| return result |
| elif isinstance(a, variables.ConstantVariable) and isinstance( |
| b, variables.ConstantVariable |
| ): |
| if self.fn is max: |
| return variables.ConstantVariable(max(a.value, b.value)) |
| else: |
| return variables.ConstantVariable(min(a.value, b.value)) |
| elif isinstance(a, SymNodeVariable) or isinstance(b, SymNodeVariable): |
| proxy = tx.output.create_proxy( |
| "call_function", self.fn, *proxy_args_kwargs([a, b], {}) |
| ) |
| return SymNodeVariable.create(tx, proxy, None) |
| else: |
| unimplemented(f"unsupported min / max over args {str(a)}, {str(b)}") |
| |
| call_min = _call_min_max |
| call_max = _call_min_max |
| |
| def call_range(self, tx, *args): |
| if self.unspec_python_args(*args) or self.constant_args(*args): |
| args, _ = specialize_args_kwargs(tx, args, {}) |
| return variables.RangeVariable(args) |
| elif self._dynamic_args(*args): |
| args = [variables.ConstantVariable(guard_if_dyn(arg)) for arg in args] |
| return variables.RangeVariable(args) |
| # None no-ops this handler and lets the driving function proceed |
| return None |
| |
| def _dynamic_args(self, *args, **kwargs): |
| return any(isinstance(x, SymNodeVariable) for x in args) or any( |
| isinstance(x, SymNodeVariable) for x in kwargs.values() |
| ) |
| |
| def call_slice(self, tx, *args): |
| return variables.SliceVariable(args) |
| |
| def _dyn_proxy(self, tx, *args, **kwargs): |
| from .builder import wrap_fx_proxy |
| |
| options = VariableTracker.propagate(self, args, kwargs.values()) |
| return wrap_fx_proxy( |
| tx, |
| tx.output.create_proxy( |
| "call_function", self.fn, *proxy_args_kwargs(args, kwargs) |
| ), |
| **options, |
| ) |
| |
| def _call_iter_tuple_list(self, tx, obj=None, *args, **kwargs): |
| if self._dynamic_args(*args, **kwargs): |
| return self._dyn_proxy(tx, *args, **kwargs) |
| cls = variables.BaseListVariable.cls_for(self.fn) |
| if obj is None: |
| if cls is SetVariable: |
| return cls( |
| tx, |
| [], |
| mutable_local=MutableLocal(), |
| ) |
| else: |
| return cls( |
| [], |
| mutable_local=MutableLocal(), |
| ) |
| elif obj.has_unpack_var_sequence(tx): |
| guards = set() |
| if obj.source and not is_constant_source(obj.source): |
| if isinstance(obj, TupleIteratorVariable): |
| guards.add(obj.source.make_guard(GuardBuilder.TUPLE_ITERATOR_LEN)) |
| else: |
| guards.add(obj.source.make_guard(GuardBuilder.LIST_LENGTH)) |
| if cls is SetVariable: |
| return cls( |
| tx, |
| list(obj.unpack_var_sequence(tx)), |
| mutable_local=MutableLocal(), |
| guards=guards, |
| ).add_options(self, obj) |
| |
| return cls( |
| list(obj.unpack_var_sequence(tx)), |
| mutable_local=MutableLocal(), |
| guards=guards, |
| ).add_options(self, obj) |
| |
| call_iter = _call_iter_tuple_list |
| call_tuple = _call_iter_tuple_list |
| call_list = _call_iter_tuple_list |
| call_set = _call_iter_tuple_list |
| |
| @staticmethod |
| def is_supported_call_dict_arg(tx, arg): |
| return ( |
| arg is None |
| or isinstance(arg, ConstDictVariable) |
| or ( |
| isinstance( |
| arg, |
| ( |
| ListVariable, |
| TupleVariable, |
| ListIteratorVariable, |
| ), |
| ) |
| and all( |
| isinstance(x, (ListVariable, TupleVariable)) |
| and isinstance( |
| x.unpack_var_sequence(tx)[0], (ConstantVariable, EnumVariable) |
| ) |
| for x in arg.unpack_var_sequence(tx) |
| ) |
| ) |
| ) |
| |
| def call_callable(self, tx, arg): |
| from .functions import BaseUserFunctionVariable |
| |
| if isinstance( |
| arg, (variables.UserDefinedClassVariable, BaseUserFunctionVariable) |
| ): |
| return variables.ConstantVariable(True).add_options(arg) |
| |
| @staticmethod |
| def call_dict_helper(tx, user_cls, arg, **options): |
| if arg is None or isinstance(arg, dict): |
| return ConstDictVariable( |
| arg if arg is not None else {}, user_cls, mutable_local=MutableLocal() |
| ).add_options(options) |
| elif isinstance(arg, variables.ConstDictVariable): |
| return arg.clone( |
| user_cls=user_cls, mutable_local=MutableLocal() |
| ).add_options(options) |
| elif isinstance( |
| arg, |
| ( |
| ListVariable, |
| TupleVariable, |
| ListIteratorVariable, |
| ), |
| ): |
| items = user_cls() |
| for x in arg.unpack_var_sequence(tx): |
| k = x.unpack_var_sequence(tx)[0].as_python_constant() |
| v = x.unpack_var_sequence(tx)[1] |
| items.update({k: v}) |
| return ConstDictVariable( |
| items, user_cls, mutable_local=MutableLocal() |
| ).add_options(options) |
| else: |
| raise AssertionError("call_dict_helper with illegal arg") |
| |
| def call_cast(self, _, *args, **kwargs): |
| if len(args) == 2: |
| return args[1] |
| |
| unimplemented(f"unsupported args to builtin cast(): {args} {kwargs}") |
| |
| def call_dict(self, tx, *args, **kwargs): |
| if not (args or kwargs): |
| return self.call_dict_helper(tx, dict, None) |
| elif ( |
| not kwargs |
| and len(args) == 1 |
| and self.is_supported_call_dict_arg(tx, args[0]) |
| ): |
| return self.call_dict_helper(tx, dict, args[0]) |
| elif not args and kwargs: |
| return variables.ConstDictVariable( |
| dict(kwargs), user_cls=dict, mutable_local=MutableLocal() |
| ) |
| else: |
| unimplemented(f"dict(): {args} {kwargs}") |
| |
| def call_zip(self, tx, *args): |
| options = VariableTracker.propagate(self, args) |
| if all(x.has_unpack_var_sequence(tx) for x in args): |
| items = [ |
| variables.TupleVariable(list(item), **options) |
| for item in zip(*[arg.unpack_var_sequence(tx) for arg in args]) |
| ] |
| return variables.TupleVariable(items, **options) |
| |
| def call_enumerate(self, tx, *args): |
| options = VariableTracker.propagate(self, args) |
| if len(args) == 1: |
| start = 0 |
| else: |
| assert len(args) == 2 |
| assert isinstance(args[1], variables.ConstantVariable) |
| start = args[1].as_python_constant() |
| if args[0].has_unpack_var_sequence(tx): |
| items = [ |
| variables.TupleVariable( |
| [variables.ConstantVariable(idx, **options), var], |
| **options, |
| ) |
| for idx, var in enumerate(args[0].unpack_var_sequence(tx), start) |
| ] |
| return variables.TupleVariable(items, **options) |
| |
| def call_len(self, tx, *args, **kwargs): |
| return args[0].call_method(tx, "__len__", args[1:], kwargs) |
| |
| def call_getitem(self, tx, *args, **kwargs): |
| if self.unspec_python_args(*args, **kwargs): |
| args, kwargs = specialize_args_kwargs(tx, args, kwargs) |
| return args[0].call_method(tx, "__getitem__", args[1:], kwargs) |
| |
| def call_isinstance(self, tx, arg, isinstance_type): |
| arg_type = arg.python_type() |
| |
| isinstance_type = isinstance_type.as_python_constant() |
| |
| if isinstance(arg, variables.TensorVariable) and arg.dtype is not None: |
| return variables.ConstantVariable(arg.call_isinstance(isinstance_type)) |
| # UserDefinedObject with C extensions can have torch.Tensor attributes, |
| # so break graph. |
| if isinstance(arg, variables.UserDefinedObjectVariable) and isinstance( |
| arg.value, types.MemberDescriptorType |
| ): |
| unimplemented( |
| f"isinstance called on UserDefinedClass {arg} {isinstance_type}" |
| ) |
| # handle __instancecheck__ defined in user class |
| if ( |
| isinstance(arg, variables.UserDefinedObjectVariable) |
| and "__instancecheck__" in isinstance_type.__class__.__dict__ |
| ): |
| return variables.ConstantVariable( |
| isinstance_type.__class__.__instancecheck__(isinstance_type, arg.value) |
| ) |
| |
| try: |
| val = issubclass(arg_type, isinstance_type) |
| except TypeError: |
| val = arg_type is isinstance_type |
| return variables.ConstantVariable(val) |
| |
| def call_super(self, tx, a, b): |
| source = ( |
| None |
| if a.source is None or b.source is None |
| else SuperSource(type=a.source, base=b.source) |
| ) |
| return variables.SuperVariable(a, b, source=source) |
| |
| def call_next(self, tx, arg): |
| if isinstance(arg, variables.ListIteratorVariable): |
| val, next_iter = arg.next_variables() |
| tx.replace_all(arg, next_iter) |
| return val |
| elif isinstance(arg, variables.BaseListVariable): |
| return arg.items[0].add_options(self, arg) |
| |
| def call_hasattr(self, tx, obj, attr): |
| if attr.is_python_constant(): |
| name = attr.as_python_constant() |
| return obj.call_hasattr(tx, name).add_options(self, obj, attr) |
| |
| def call_map(self, tx, fn, seq): |
| if seq.has_unpack_var_sequence(tx): |
| items = [fn.call_function(tx, [x], {}) for x in seq.unpack_var_sequence(tx)] |
| return variables.TupleVariable(items).add_options(self, fn, seq) |
| |
| def call_sum(self, tx, seq, **kwargs): |
| # Special case for sum on tuple of floats and ints |
| if ( |
| isinstance(seq, (variables.ListVariable, variables.TupleVariable)) |
| and all( |
| isinstance(x, variables.ConstantVariable) |
| and isinstance(x.value, (int, float)) |
| for x in seq.items |
| ) |
| and not kwargs |
| ): |
| new_list = [x.value for x in seq.items] |
| return variables.ConstantVariable(sum(new_list)) |
| if seq.has_unpack_var_sequence(tx): |
| start = kwargs.pop( |
| "start", variables.ConstantVariable(0) |
| ).as_python_constant() |
| assert not kwargs |
| items = seq.unpack_var_sequence(tx)[start:] |
| return BuiltinVariable(functools.reduce).call_function( |
| tx, |
| [ |
| BuiltinVariable(operator.add), |
| variables.TupleVariable(items), |
| variables.ConstantVariable(0).add_options(self, seq), |
| ], |
| {}, |
| ) |
| |
| def call_reduce(self, tx, function, iterable, initializer=None): |
| if iterable.has_unpack_var_sequence(tx): |
| items = iterable.unpack_var_sequence(tx) |
| if initializer is None: |
| value, items = items[0], items[1:] |
| else: |
| value = initializer |
| for element in items: |
| value = function.call_function(tx, [value, element], {}) |
| return value |
| |
| def call_getattr( |
| self, tx, obj: VariableTracker, name_var: VariableTracker, default=None |
| ): |
| from . import ( |
| ConstantVariable, |
| GetAttrVariable, |
| PythonModuleVariable, |
| TorchVariable, |
| UserFunctionVariable, |
| ) |
| from .builder import VariableBuilder |
| |
| options = VariableTracker.propagate(self, obj, name_var) |
| guards = options["guards"] |
| name = name_var.as_python_constant() |
| |
| if not name_var.is_python_constant(): |
| unimplemented("non-const getattr() name") |
| |
| if tx.output.side_effects.is_attribute_mutation(obj): |
| try: |
| # re-read a pending side effect? |
| return tx.output.side_effects.load_attr(obj, name).add_options(options) |
| except KeyError: |
| pass |
| |
| if default is not None: |
| hasattr_var = self.call_hasattr(tx, obj, name_var) |
| guards.update(hasattr_var.guards) |
| assert hasattr_var.as_python_constant() in (True, False) |
| if not hasattr_var.as_python_constant(): |
| return default.add_guards(guards) |
| |
| if obj.source: |
| source = AttrSource(obj.source, name) |
| options["source"] = source |
| else: |
| source = None |
| |
| if isinstance(obj, variables.NNModuleVariable): |
| return obj.var_getattr(tx, name).add_options(options) |
| elif isinstance(obj, variables.TensorVariable) and name == "grad": |
| if source: |
| # We are going to be raising this tensor as grapharg. So, ensure |
| # that we have real grad value instead of fake tensor value. |
| # Walk through the inputs of the subgraph and find if we already |
| # have the original tensor stored in the graphargs. |
| for grapharg in tx.output.graphargs: |
| if grapharg.source == source.base: |
| example_value = grapharg.example.grad |
| return VariableBuilder(tx, source)(example_value).add_options( |
| options |
| ) |
| unimplemented("tensor grad") |
| else: |
| unimplemented("tensor grad") |
| elif isinstance( |
| obj, |
| ( |
| variables.TensorVariable, |
| variables.NamedTupleVariable, |
| variables.ConstantVariable, |
| variables.UserDefinedClassVariable, |
| variables.UserDefinedObjectVariable, |
| ), |
| ): |
| try: |
| return ( |
| obj.var_getattr(tx, name).clone(source=source).add_options(options) |
| ) |
| except NotImplementedError: |
| return GetAttrVariable(obj, name, **options) |
| elif isinstance(obj, TorchVariable): |
| member = getattr(obj.value, name) |
| if is_utils_checkpoint(member): |
| options["source"] = source |
| return build_checkpoint_variable(**options) |
| elif is_allowed(member): |
| return TorchVariable(member, **options) |
| elif ConstantVariable.is_literal(member): |
| return ConstantVariable(member, **options) |
| else: |
| return VariableBuilder(tx, source)(member).add_guards(guards) |
| elif isinstance(obj, (PythonModuleVariable, DummyModule)): |
| member = obj.value.__dict__[name] |
| |
| if config.replay_record_enabled: |
| tx.exec_recorder.record_module_access(obj.value, name, member) |
| |
| return VariableBuilder(tx, source)(member).add_guards(guards) |
| elif istype(obj, UserFunctionVariable) and name in ("__name__", "__module__"): |
| return ConstantVariable( |
| getattr(obj.fn, name), **VariableTracker.propagate(obj) |
| ) |
| else: |
| try: |
| return ( |
| obj.var_getattr(tx, name).clone(source=source).add_options(options) |
| ) |
| except NotImplementedError: |
| return GetAttrVariable(obj, name, **options) |
| |
| def call_setattr( |
| self, tx, obj: VariableTracker, name_var: VariableTracker, val: VariableTracker |
| ): |
| from .distributed import PlacementVariable |
| |
| if isinstance(obj, (variables.DataClassVariable, PlacementVariable)): |
| return obj.call_method(tx, "__setattr__", [name_var, val], {}) |
| elif ( |
| tx.output.side_effects.is_attribute_mutation(obj) |
| and name_var.is_python_constant() |
| ): |
| tx.output.side_effects.store_attr(obj, name_var.as_python_constant(), val) |
| return val.add_options(self, obj, name_var) |
| elif isinstance(obj, variables.UserDefinedObjectVariable): |
| unimplemented( |
| f"setattr(UserDefinedObjectVariable) {type(obj.value).__setattr__}" |
| ) |
| elif isinstance(obj, variables.NNModuleVariable): |
| if not tx.output.is_root_tracer(): |
| raise AttributeMutationError( |
| "Can't inplace modify module params/buffers inside HigherOrderOp" |
| ) |
| if name_var.is_python_constant() and isinstance( |
| val, variables.TensorVariable |
| ): |
| assigning_fake_val = get_fake_value(val.as_proxy().node, tx) |
| |
| try: |
| getattr_var = obj.var_getattr(tx, name_var.as_python_constant()) |
| except AttributeError: |
| getattr_var = None |
| |
| if isinstance(getattr_var, variables.TensorVariable): |
| # get_fake_val will return a real tensor here because it's an attribute on the module (get_attr node) |
| existing_attr = get_fake_value(getattr_var.as_proxy().node, tx) |
| existing_fake_attr = ( |
| variables.builder.wrap_to_fake_tensor_and_record( |
| existing_attr, tx, source=getattr_var.source, is_tensor=True |
| ) |
| ) |
| |
| # same tensor identiy, setattr is a no-op |
| mod_setattr = inspect.getattr_static(obj.module_type, "__setattr__") |
| if ( |
| existing_fake_attr is assigning_fake_val |
| and mod_setattr is torch.nn.Module.__setattr__ |
| ): |
| return getattr_var |
| |
| obj.convert_to_unspecialized(tx) |
| |
| def call_delattr(self, tx, obj: VariableTracker, name_var: VariableTracker): |
| return self.call_setattr(tx, obj, name_var, variables.DeletedVariable()) |
| |
| def call_type(self, tx, obj: VariableTracker): |
| from .builder import VariableBuilder |
| |
| try: |
| py_type = obj.python_type() |
| except NotImplementedError: |
| py_type = None |
| |
| if istype(obj, variables.TupleVariable): |
| return BuiltinVariable(py_type).add_options(self, obj) |
| |
| if py_type is not None and obj.source: |
| return VariableBuilder(tx, TypeSource(obj.source))(py_type).add_options( |
| self, obj |
| ) |
| |
| raise UserError( |
| UserErrorType.ANTI_PATTERN, |
| "Can't call type() on generated custom object. " |
| "Please use __class__ instead", |
| ) |
| |
| def call_reversed(self, tx, obj: VariableTracker): |
| if obj.has_unpack_var_sequence(tx): |
| items = list(reversed(obj.unpack_var_sequence(tx))) |
| return variables.TupleVariable( |
| items, **VariableTracker.propagate(self, obj) |
| ) |
| |
| def call_sorted(self, tx, obj: VariableTracker, **kwargs): |
| if ( |
| obj.has_unpack_var_sequence(tx) |
| and not isinstance(obj, variables.TensorVariable) |
| and all(x.is_python_constant() for x in obj.unpack_var_sequence(tx)) |
| ): |
| function = kwargs.pop("key", None) |
| reverse = kwargs.pop( |
| "reverse", ConstantVariable(False) |
| ).as_python_constant() |
| assert len(kwargs) == 0 |
| if function: |
| items = sorted( |
| obj.unpack_var_sequence(tx), |
| key=lambda x: function.call_function( |
| tx, [x], {} |
| ).as_python_constant(), |
| reverse=reverse, |
| ) |
| else: |
| items = sorted( |
| obj.unpack_var_sequence(tx), |
| key=lambda x: x.as_python_constant(), |
| reverse=reverse, |
| ) |
| return variables.ListVariable(items, **VariableTracker.propagate(self, obj)) |
| |
| def call_chain(self, tx, *args): |
| if all(obj.has_unpack_var_sequence(tx) for obj in args): |
| items = [] |
| for obj in args: |
| items.extend(obj.unpack_var_sequence(tx)) |
| return variables.TupleVariable( |
| items, **VariableTracker.propagate(self, *args) |
| ) |
| |
| def call_islice(self, tx, iterable, *args): |
| if iterable.has_unpack_var_sequence(tx) and all( |
| x.is_python_constant() for x in args |
| ): |
| const_args = [x.as_python_constant() for x in args] |
| items = iterable.unpack_var_sequence(tx) |
| items = list(itertools.islice(items, *const_args)) |
| return variables.TupleVariable( |
| items, **VariableTracker.propagate(self, iterable, *args) |
| ) |
| |
| # neg is a constant fold function, so we only get here if constant fold is not valid |
| def call_neg(self, tx, a): |
| if isinstance(a, SymNodeVariable): |
| return SymNodeVariable.create( |
| tx, |
| (operator.neg)(a.as_proxy()), |
| sym_num=None, |
| ) |
| # None no-ops this handler and lets the driving function proceed |
| return None |
| |
| def call_id(self, tx, *args): |
| if len(args) > 0 and isinstance(args[0], variables.NNModuleVariable): |
| nn_mod_variable = args[0] |
| mod = tx.output.get_submodule(nn_mod_variable.module_key) |
| return variables.ConstantVariable(id(mod)) |
| else: |
| unimplemented(f"call_id with args {args}") |
| |
| def _comparison(self, tx, left, right): |
| """ |
| Used to implement comparison operators for different types. |
| For example, list1 < list2 is implemented differently from tensor1 < tensor2 |
| """ |
| from . import ( |
| BaseListVariable, |
| ConstantVariable, |
| NNModuleVariable, |
| TensorVariable, |
| UserDefinedObjectVariable, |
| UserFunctionVariable, |
| ) |
| from .lists import SizeVariable |
| from .tensor import ( |
| supported_const_comparison_ops, |
| supported_tensor_comparison_ops, |
| ) |
| |
| op = self.fn |
| |
| def _unimplemented(): |
| unimplemented(f"comparison {typestr(left)} {op} {typestr(right)}") |
| |
| if ( |
| all( |
| isinstance(x, (NNModuleVariable, ConstantVariable)) |
| for x in [left, right] |
| ) |
| and op in supported_const_comparison_ops.values() |
| ): |
| left = ( |
| tx.output.get_submodule(left.module_key) |
| if isinstance(left, NNModuleVariable) |
| else left.as_python_constant() |
| ) |
| right = ( |
| tx.output.get_submodule(right.module_key) |
| if isinstance(right, NNModuleVariable) |
| else right.as_python_constant() |
| ) |
| return ConstantVariable(op(left, right)) |
| |
| if isinstance(left, UserFunctionVariable): |
| if op not in supported_const_comparison_ops.values(): |
| _unimplemented() |
| if not isinstance(right, UserFunctionVariable): |
| _unimplemented() |
| return ConstantVariable(op(left.fn, right.fn)) |
| |
| # Note, we have a rare BaseListVariable subtype mismatch with valid comparison |
| # x = torch.randn([3, 3]) |
| # x.size() == (3, 3) # True |
| # (3, 3) == x.size() # True |
| if isinstance(left, (SizeVariable, TupleVariable)) and isinstance( |
| right, (TupleVariable, SizeVariable) |
| ): |
| return BaseListVariable.list_compare(tx, op, left, right) |
| |
| if isinstance(left, BaseListVariable): |
| if not type(left) == type(right): # Mismatch in BaseListVariable subclasses |
| _unimplemented() |
| return BaseListVariable.list_compare(tx, op, left, right) |
| |
| if isinstance(left, SetVariable): |
| if not type(left) == type(right): # Mismatch in BaseListVariable subclasses |
| _unimplemented() |
| return ConstantVariable(op(left._underlying_items, right._underlying_items)) |
| |
| if isinstance(left, TensorVariable): |
| from .builder import wrap_fx_proxy_cls |
| |
| if op not in supported_tensor_comparison_ops.values(): |
| _unimplemented() |
| return wrap_fx_proxy_cls( |
| type(left), # handle Ndarrays and Tensors |
| tx, |
| op(left.as_proxy(), right.as_proxy()), |
| ) |
| |
| if isinstance(left, SymNodeVariable) or isinstance(right, SymNodeVariable): |
| if op not in supported_tensor_comparison_ops.values(): |
| _unimplemented() |
| |
| return SymNodeVariable.create( |
| tx, |
| op(left.as_proxy(), right.as_proxy()), |
| sym_num=None, |
| ) |
| |
| if isinstance(left, ConstantVariable) and isinstance(right, ConstantVariable): |
| return ConstantVariable(op(left.value, right.value)) |
| |
| if isinstance(left, UserDefinedObjectVariable) and isinstance( |
| right, UserDefinedObjectVariable |
| ): |
| return ConstantVariable(op(left.value, right.value)) |
| |
| if op.__name__ == "is_": |
| # If the two objects are of different type, we can safely return False |
| if type(left) is not type(right): |
| return ConstantVariable(False) |
| |
| _unimplemented() |
| |
| # and_ is a constant fold function, so we only get here if constant fold is not valid |
| def call_and_(self, tx, a, b): |
| if isinstance(a, (SymNodeVariable, ConstantVariable)) and isinstance( |
| b, (SymNodeVariable, ConstantVariable) |
| ): |
| return SymNodeVariable.create( |
| tx, |
| tx.output.create_proxy( |
| "call_function", operator.and_, *proxy_args_kwargs([a, b], {}) |
| ), |
| sym_num=None, |
| ) |
| # None no-ops this handler and lets the driving function proceed |
| return None |
| |
| # or_ is a constant fold function, so we only get here if constant fold is not valid |
| def call_or_(self, tx, a, b): |
| if isinstance(a, (SymNodeVariable, ConstantVariable)) and isinstance( |
| b, (SymNodeVariable, ConstantVariable) |
| ): |
| return SymNodeVariable.create( |
| tx, |
| tx.output.create_proxy( |
| "call_function", operator.or_, *proxy_args_kwargs([a, b], {}) |
| ), |
| sym_num=None, |
| ) |
| # None no-ops this handler and lets the driving function proceed |
| return None |
| |
| def call_not_(self, tx, a): |
| if isinstance(a, SymNodeVariable): |
| return SymNodeVariable.create( |
| tx, |
| tx.output.create_proxy( |
| "call_function", operator.not_, *proxy_args_kwargs([a], {}) |
| ), |
| sym_num=None, |
| ) |
| |
| if isinstance(a, ListVariable): |
| return ConstantVariable(len(a.items) == 0).add_options(self, a) |
| |
| return None |
| |
| call_eq = _comparison |
| call_gt = _comparison |
| call_lt = _comparison |
| call_ge = _comparison |
| call_le = _comparison |
| call_ne = _comparison |
| call_is_ = _comparison |
| call_is_not = _comparison |
