[dynamo][user-defined] Simplify and improve scope of UserDefinedObject var_getattr (pytorch#130169)

Fixes pytorch#122649

Pull Request resolved: pytorch#130169
Approved by: https://github.com/jansel
ghstack dependencies: pytorch#118448, pytorch#130159

Author: anijain2305

test/dynamo/test_export.py

@@ -3428,8 +3428,7 @@ def f_mismatch_return_length(x):
 
         example_inputs = (torch.rand(5),)
         with self.assertRaisesRegex(
-            torch._dynamo.exc.UncapturedHigherOrderOpError,
-            "Cond doesn't work unless it is captured completely with torch.compile",
+            RuntimeError, "Unmatched number of outputs from cond"
         ):
             torch._dynamo.export(
                 f_mismatch_return_length,

test/dynamo_expected_failures/TestQuantizeFx.test_prepare_custom_config_set_standalone_module_class renamed to test/dynamo_expected_failures/TestSerialization.test_serialization_dill

@@ -1411,8 +1411,8 @@ def f(x, y):
 
         x = torch.randn(4)
         with self.assertRaisesRegex(
-            torch._dynamo.exc.UncapturedHigherOrderOpError,
-            "Cond doesn't work unless it is captured completely with torch.compile",
+            torch._dynamo.exc.CondOpArgsMismatchError,
+            "Expected to return same number of outputs but got:",
         ):
             make_fx(f)(x, torch.tensor(False))
 
@@ -1584,8 +1584,8 @@ def f(x, y):
 
         x = torch.randn(4)
         with self.assertRaisesRegex(
-            torch._dynamo.exc.UncapturedHigherOrderOpError,
-            "Cond doesn't work unless it is captured completely with torch.compile",
+            torch._dynamo.exc.CondOpArgsMismatchError,
+            "Expected to return same number of outputs but got:",
         ):
             make_fx(f, tracing_mode="fake")(x, torch.tensor(False))
 

test/dynamo_expected_failures/TestTensorProtoSummary.test_half_tensor_proto_bfloat16_proto_type_14

@@ -45,6 +45,7 @@
     run_tests,
     TEST_WITH_CROSSREF,
     TestCase,
+    skipIfTorchDynamo,
 )
 
 
@@ -760,6 +761,7 @@ class TestTensorProtoSummary(BaseTestCase):
             (torch.bfloat16, DataType.DT_BFLOAT16),
         ],
     )
+    @skipIfTorchDynamo("Unsuitable test for Dynamo, behavior changes with version")
     def test_half_tensor_proto(self, tensor_type, proto_type):
         float_values = [1.0, 2.0, 3.0]
         actual_proto = tensor_proto(

test/dynamo_expected_failures/TestTensorProtoSummary.test_half_tensor_proto_float16_proto_type_19

@@ -587,9 +587,16 @@ def is_wrapper_or_member_descriptor(value):
     return isinstance(
         value,
         (
-            types.MethodWrapperType,
+            # set up by PyGetSetDef
+            types.GetSetDescriptorType,
+            # set by PyMethodDef, e.g. list.append
+            types.MethodDescriptorType,
+            # slots - list.__add__
             types.WrapperDescriptorType,
+            # set up by PyMemberDef
             types.MemberDescriptorType,
+            # wrapper over C functions
+            types.MethodWrapperType,
         ),
     )
 

test/functorch/test_control_flow.py

@@ -2538,6 +2538,22 @@ def make_type_handlers():
         handlers[immutable_list] = handlers[list]
         handlers[types.ModuleType] = lambda tx, value: PythonModuleVariable(value)
 
+        handlers[
+            torch.distributions.constraints._Real
+        ] = lambda tx, value: UserDefinedObjectVariable(
+            value, mutable_local=MutableLocal()
+        )
+        handlers[
+            torch.distributions.constraints._Interval
+        ] = lambda tx, value: UserDefinedObjectVariable(
+            value, mutable_local=MutableLocal()
+        )
+        handlers[
+            torch.distributions.constraints.Constraint
+        ] = lambda tx, value: UserDefinedObjectVariable(
+            value, mutable_local=MutableLocal()
+        )
+
         def passthrough(tx, value):
             return value
 

test/test_tensorboard.py

@@ -4,11 +4,9 @@
 import contextlib
 import enum
 import functools
-import importlib
 import inspect
 import itertools
 import random
-import re
 import sys
 import threading
 import types
@@ -45,7 +43,6 @@
     WeakRefCallSource,
 )
 from ..utils import (
-    all_hook_names,
     build_checkpoint_variable,
     check_constant_args,
     get_custom_getattr,
@@ -830,7 +827,6 @@ def is_supported_nn_module_method(self, method):
     def var_getattr(self, tx, name):
         from .. import trace_rules
         from . import ConstantVariable
-        from .builder import VariableBuilder
 
         value = self.value
         source = AttrSource(self.source, name) if self.source else None
@@ -843,6 +839,13 @@ def var_getattr(self, tx, name):
             options = {"source": source}
             return variables.GetAttrVariable(self, name, **options)
 
+        # TODO(anijain2305) - Investigate if we need specialization for more
+        # dunder attrs. inspect.getattr_static does not return correct value for
+        # them.
+        if name == "__class__":
+            options = {"source": source}
+            return UserDefinedClassVariable(type(self.value), **options)
+
         try:
             subobj = self._getattr_static(name)
         except AttributeError:
@@ -958,75 +961,15 @@ def var_getattr(self, tx, name):
                 else:
                     return trace_rules.lookup(func)(func)
 
-        if (
-            name in getattr(value, "__dict__", {})
-            or ConstantVariable.is_literal(subobj)
-            or isinstance(
-                subobj,
-                (
-                    torch.Tensor,
-                    torch.nn.Module,
-                    re.Pattern,
-                ),
-            )
-        ):
+        if subobj is not NO_SUCH_SUBOBJ and not is_wrapper_or_member_descriptor(subobj):
             if source:
                 return variables.LazyVariableTracker.create(subobj, source)
-            elif ConstantVariable.is_literal(subobj):
-                return ConstantVariable.create(subobj)
-            elif (
-                type(subobj) == torch.utils._pytree.TreeSpec
-                or type(subobj) == torch.utils._pytree.LeafSpec
-                or type(value) == torch.utils._pytree.TreeSpec
-            ):
+            else:
                 from .builder import SourcelessBuilder
 
                 return SourcelessBuilder.create(tx, subobj)
 
-        if (
-            subobj is not NO_SUCH_SUBOBJ
-            and name not in getattr(value, "__dict__", {})
-            and (
-                type(value).__module__.startswith("torch.")
-                or isinstance(subobj, re.Pattern)
-            )
-            and "torch.optim" not in type(value).__module__
-            and not callable(value)
-            and not isinstance(subobj, types.MethodDescriptorType)
-        ):
-            if not source:
-                assert getattr(
-                    importlib.import_module(type(value).__module__),
-                    type(value).__name__,
-                ) is type(value)
-                source = AttrSource(
-                    AttrSource(
-                        tx.import_source(type(value).__module__), type(value).__name__
-                    ),
-                    name,
-                )
-
-            return VariableBuilder(tx, source)(subobj)
         options = {"source": source}
-        if isinstance(
-            subobj,
-            (
-                torch.distributions.constraints._Interval,
-                torch.distributions.constraints._Real,
-                torch.distributions.constraints.Constraint,
-            ),
-        ):
-            return UserDefinedObjectVariable(subobj, **options)
-        elif isinstance(self.value, torch.nn.Module) and name in all_hook_names:
-            assert isinstance(subobj, collections.OrderedDict)
-            if not subobj:
-                return variables.ConstDictVariable(
-                    subobj, collections.OrderedDict, **options
-                )
-
-        if name == "__class__":
-            return UserDefinedClassVariable(type(self.value), **options)
-
         return variables.GetAttrVariable(self, name, **options)
 
     def call_hasattr(self, tx, name: str) -> "VariableTracker":