[JIT] Add property support for ScriptModules

Meghan Lele · Meghan Lele · commit 0e428f40d92e · 2020-08-14T19:05:12.000-07:00
**Summary** This commit extends support for properties to include ScriptModules. **Test Plan** This commit adds a unit test that has a ScriptModule with a user-defined property. `python test/test_jit_py3.py TestScriptPy3.test_module_properties` ghstack-source-id: 4d97cd7 Pull Request resolved: #42390
diff --git a/test/jit/test_recursive_script.py b/test/jit/test_recursive_script.py
@@ -625,24 +625,6 @@ def forward(self, x):
         m = M()
         self.checkModule(m, (torch.randn(5, 5), ))
 
-    def test_property(self):
-        class M(nn.Module):
-            def __init__(self):
-                super(M, self).__init__()
-                self.x = 0
-
-            @property
-            def x_and_1(self):
-                return self.x + 1
-
-            def forward(self, new_x):
-                # type: (int) -> int
-                self.x = new_x
-                return self.x_and_1
-
-        with self.assertRaisesRegex(RuntimeError, "property"):
-            torch.jit.script(M())
-
     def test_inner_traced_module(self):
         class Dummy(nn.Module):
             def forward(self, x):
diff --git a/test/test_jit_py3.py b/test/test_jit_py3.py
@@ -533,6 +533,55 @@ def forward(self, x):
         # Check that ignored method is still intact.
         self.assertEqual(inp, n.ignored_method(inp))
 
+    def test_module_properties(self):
+        class ModuleWithProperties(torch.nn.Module):
+            def __init__(self, a: int):
+                super().__init__()
+                self.a = a
+
+            def forward(self, a: int, b: int):
+                self.attr = a + b
+                return self.attr
+
+            @property
+            def attr(self):
+                return self.a
+
+            @torch.jit.ignore
+            @property
+            def ignored_attr(self):
+                return sum([self.a])
+
+            @attr.setter
+            def attr(self, a: int):
+                if a > 0:
+                    self.a = a
+                else:
+                    self.a = 0
+
+        class ModuleWithNoSetter(torch.nn.Module):
+            def __init__(self, a: int):
+                super().__init__()
+                self.a = a
+
+            def forward(self, a: int, b: int):
+                self.attr + a + b
+
+            @property
+            def attr(self):
+                return self.a + 1
+
+        self.checkModule(ModuleWithProperties(5), (5, 6,))
+        self.checkModule(ModuleWithProperties(5), (-5, -6,))
+        self.checkModule(ModuleWithNoSetter(5), (5, 6,))
+        self.checkModule(ModuleWithNoSetter(5), (-5, -6,))
+
+        mod = ModuleWithProperties(3)
+        scripted_mod = torch.jit.script(mod)
+
+        with self.assertRaisesRegex(torch.nn.modules.module.ModuleAttributeError, "has no attribute"):
+            scripted_mod.ignored_attr
+
     def test_export_opnames_interface(self):
         global OneTwoModule
 
diff --git a/torch/_jit_internal.py b/torch/_jit_internal.py
@@ -463,6 +463,15 @@ def forward(self, x):
         fn._torchscript_modifier = FunctionModifiers.IGNORE
         return fn
 
+    if isinstance(drop, property):
+        prop = drop
+        setattr(prop.fget, "_torchscript_modifier", FunctionModifiers.IGNORE)
+
+        if prop.fset:
+            setattr(prop.fset, "_torchscript_modifier", FunctionModifiers.IGNORE)
+
+        return prop
+
     if not isinstance(drop, bool):
         raise RuntimeError("Argument to @torch.jit.ignore must be a bool or "
                            "a function but got {}".format(drop))
diff --git a/torch/csrc/jit/python/python_sugared_value.cpp b/torch/csrc/jit/python/python_sugared_value.cpp
@@ -599,6 +599,14 @@ std::shared_ptr<SugaredValue> ModuleValue::attr(
     return attr;
   }
 
+  // Check if it's a property.
+  auto prop =
+      concreteType_->getJitType()->expect<ClassType>()->getProperty(field);
+  if (prop) {
+    return MethodValue(self_, prop->getter->name())
+        .call(loc, m, {}, {}, /*n_binders=*/1);
+  }
+
   // We don't define this attr. Bailout with a hint to the user.
   std::string hint;
   if (auto failureReason = concreteType_->findFailedAttribute(field)) {
diff --git a/torch/csrc/jit/python/python_tree_views.cpp b/torch/csrc/jit/python/python_tree_views.cpp
@@ -4,6 +4,7 @@
 
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
+#include <torch/csrc/utils/pybind.h>
 
 #include <sstream>
 
@@ -162,7 +163,18 @@ void initTreeViewBindings(PyObject* module) {
                        const Def& getter,
                        Def* setter) {
         return Property::create(r, name, getter, wrap_maybe(r, setter));
-      }));
+      }))
+      .def("name", [](const Property& property) { return property.name(); })
+      .def(
+          "getter_name",
+          [](const Property& property) { return property.getter().name(); })
+      .def("setter_name", [](const Property& property) {
+        if (property.setter().present()) {
+          return c10::optional<Ident>(property.setter().get().name());
+        }
+
+        return c10::optional<Ident>(c10::nullopt);
+      });
 
   py::class_<ClassDef, TreeView>(m, "ClassDef")
       .def(py::init([](const Ident& name,
diff --git a/torch/csrc/jit/python/script_init.cpp b/torch/csrc/jit/python/script_init.cpp
@@ -1569,33 +1569,43 @@ void initJitScriptBindings(PyObject* module) {
             return self.equals(other);
           })
       .def(
-          "_create_methods",
+          "_create_methods_and_properties",
           [](std::shared_ptr<ConcreteModuleType> concreteType,
-             const std::vector<Def>& defs,
-             const std::vector<ResolutionCallback>& rcbs,
+             const std::vector<Property>& properties,
+             const std::vector<ResolutionCallback>& propertyRcbs,
+             const std::vector<Def>& methodDefs,
+             const std::vector<ResolutionCallback>& methodRcbs,
              const std::vector<FunctionDefaults>& defaults) {
-            TORCH_INTERNAL_ASSERT(defs.size() == rcbs.size());
-            std::vector<ResolverPtr> resolvers;
-            resolvers.reserve(rcbs.size());
-            for (auto& callback : rcbs) {
-              resolvers.push_back(pythonResolver(callback));
+            TORCH_INTERNAL_ASSERT(methodDefs.size() == methodRcbs.size());
+            TORCH_INTERNAL_ASSERT(properties.size() == propertyRcbs.size());
+
+            std::vector<ResolverPtr> methodResolvers, propertyResolvers;
+            methodResolvers.reserve(methodRcbs.size());
+            for (auto& callback : methodRcbs) {
+              methodResolvers.push_back(pythonResolver(callback));
+            }
+
+            propertyResolvers.reserve(propertyRcbs.size());
+            for (auto& callback : propertyRcbs) {
+              propertyResolvers.push_back(pythonResolver(callback));
             }
+
             const auto& selfType =
                 concreteType->getJitType()->expect<ClassType>();
             const auto& prefix = selfType->name().value();
             const auto self = ModuleSelf(std::move(concreteType));
             auto cu = selfType->compilation_unit();
             cu->define(
                 prefix,
-                /*properties=*/{},
-                /*propResolvers=*/{},
-                defs,
-                resolvers,
+                properties,
+                propertyResolvers,
+                methodDefs,
+                methodResolvers,
                 &self);
             // Stitch in default arguments for each Def if provided
             auto defaults_it = defaults.begin();
-            auto defs_it = defs.begin();
-            while (defs_it != defs.end()) {
+            auto defs_it = methodDefs.begin();
+            while (defs_it != methodDefs.end()) {
               const auto method_name =
                   QualifiedName(prefix, (*defs_it).name().name());
               auto& method = cu->get_function(method_name);
diff --git a/torch/jit/_recursive.py b/torch/jit/_recursive.py
@@ -6,13 +6,15 @@
 import warnings
 
 import torch._jit_internal as _jit_internal
-from torch.jit.frontend import get_default_args, get_jit_def
+from torch.jit.frontend import get_default_args, get_jit_def, get_class_properties
 from torch.jit._builtins import _find_builtin
 from torch.nn import Module
 from torch._six import get_function_from_type, bind_method
 
 
 ScriptMethodStub = collections.namedtuple('ScriptMethodStub', ('resolution_callback', 'def_', 'original_method'))
+PropertyStub = collections.namedtuple('Property', ('resolution_callback', 'def_'))
+
 
 # TODO: there should be a more principled way of doing this.
 ignored_attributes = [
@@ -48,6 +50,7 @@ def make_stub_from_method(nn_module, method_name):
     # even though we requested a stub for `forward`.
     return make_stub(func, method_name)
 
+
 # base types that can be constants
 # in addition, tuples and lists of these base types are also considered constants
 # If you edit this list, then you also need to edit the handlers in
@@ -239,14 +242,6 @@ def infer_type(name, item):
                     "to a TorchScript type.)").format(torch.typename(type(value)))
             concrete_type_builder.add_failed_attribute(name, hint)
 
-    # Add @property methods as failed attributes, to give a better error message.
-    for name, value in type(nn_module).__dict__.items():
-        if isinstance(value, property):
-            hint = ("\n(This attribute exists on the Python module, but it's an @property "
-                    "method. @property methods are not yet supported in TorchScript. "
-                    "Please file a feature request on Github)")
-            concrete_type_builder.add_failed_attribute(name, hint)
-
     return concrete_type_builder
 
 class ConcreteTypeStore(object):
@@ -285,11 +280,17 @@ def get_or_create_concrete_type(self, nn_module):
 
 concrete_type_store = ConcreteTypeStore()
 
-def create_methods_from_stubs(concrete_type, stubs):
-    defs = [m.def_ for m in stubs]
-    rcbs = [m.resolution_callback for m in stubs]
-    defaults = [get_default_args(m.original_method) for m in stubs]
-    concrete_type._create_methods(defs, rcbs, defaults)
+
+def create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs):
+    method_defs = [m.def_ for m in method_stubs]
+    method_rcbs = [m.resolution_callback for m in method_stubs]
+    method_defaults = [get_default_args(m.original_method) for m in method_stubs]
+
+    property_defs = [p.def_ for p in property_stubs]
+    property_rcbs = [p.resolution_callback for p in property_stubs]
+
+    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
+
 
 def create_script_module(nn_module, stubs_fn, share_types=True):
     """
@@ -326,7 +327,8 @@ def create_script_module_impl(nn_module, concrete_type, stubs_fn):
         stubs_fn:  Lambda that takes an nn.Module and generates a list of ScriptMethodStubs to compile.
     """
     cpp_module = torch._C._create_module_with_type(concrete_type.jit_type)
-    stubs = stubs_fn(nn_module)
+    method_stubs = stubs_fn(nn_module)
+    property_stubs = get_property_stubs(nn_module)
 
     def init_fn(script_module):
         # Initialize the ScriptModule:
@@ -354,13 +356,11 @@ def init_fn(script_module):
             cpp_module.setattr(name, scripted)
             script_module._modules[name] = scripted
 
-        # 3. Copy @ignored/@unused methods from the original `nn_module` to the new ScriptModule.
+        # 3. Copy @ignored/@unused methods and properties from the original `nn_module` to the new ScriptModule.
         #    This ensures we can access these Python methods on the ScriptModule.
         for name in dir(nn_module):
             item = getattr(nn_module, name, None)
-            if not inspect.ismethod(item):
-                continue
-            if _jit_internal.is_ignored_fn(item):
+            if inspect.ismethod(item) and _jit_internal.is_ignored_fn(item):
                 unbound_function = getattr(type(nn_module), name)
                 bound_method = unbound_function.__get__(script_module)
                 setattr(script_module, name, bound_method)
@@ -373,7 +373,7 @@ def init_fn(script_module):
 
     # Compile methods if necessary
     if concrete_type not in concrete_type_store.methods_compiled:
-        create_methods_from_stubs(concrete_type, stubs)
+        create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
         torch._C._run_emit_module_hook(cpp_module)
         concrete_type_store.methods_compiled.add(concrete_type)
 
@@ -391,14 +391,14 @@ def init_fn(script_module):
 
 
     # Make the compiled methods available to the Python ScriptModule class.
-    for stub in stubs:
-        if stub.original_method is None:
+    for method_stub in method_stubs:
+        if method_stub.original_method is None:
             # define()'d methods don't have an Python original_method, so we
             # don't need to do any Python re-wrapping stuff
             continue
 
-        name = stub.original_method.__name__
-        if name != stub.def_.name().name:
+        name = method_stub.original_method.__name__
+        if name != method_stub.def_.name().name:
             # TODO: Why skip this? Because @torch.jit._overload_method will
             # mangle the name of the function.
             continue
@@ -407,14 +407,23 @@ def init_fn(script_module):
         # Wrap the original to propagate docstrings and such.
         # TODO: we don't currently do this functions that are recursively
         # compiled, we should.
-        script_method = functools.wraps(stub.original_method)(script_method)
+        script_method = functools.wraps(method_stub.original_method)(script_method)
 
         # Add the methods to the script_module directly. This ensures they will
         # be found first when `name` is looked up (as opposed to the stubs or
         # nn.Module.forward)
         script_module.__dict__[name] = script_method
 
 
+    # Make module properties available on the Python ScriptModule class.
+    for property_stub in property_stubs:
+        property_name = property_stub.def_.name().name
+        fget = cpp_module._get_method(property_stub.def_.getter_name().name)
+        # Setter is optional, so it may not exist.
+        setter_name = property_stub.def_.setter_name()
+        fset = cpp_module._get_method(setter_name.name) if setter_name else None
+        script_module.__dict__[property_name] = property(property_name, fget, fset)
+
     # copy over python methods to script module if they aren't defined on the script module
     # this is currently an internal api used only on module containers
     for name in dir(nn_module):
@@ -548,6 +557,28 @@ def ignore_overloaded(method_name):
         stubs.append(make_stub_from_method(nn_module, method))
     return overload_stubs + stubs
 
+
+def get_property_stubs(nn_module):
+    """
+    Create property stubs for the properties of the module by creating method
+    stubs for the getter and setter.
+    """
+    module_ty = type(nn_module)
+    properties_asts = get_class_properties(module_ty, self_name="RecursiveScriptModule")
+    rcbs = {}
+
+    for name in dir(module_ty):
+        item = getattr(module_ty, name, None)
+        if isinstance(item, property):
+            if not item.fget:
+                raise RuntimeError(f'Property {name} of {nn_module.__name__} must have a getter')
+
+            rcbs[name] = _jit_internal.createResolutionCallbackFromClosure(item.fget)
+
+    stubs = [PropertyStub(rcbs[ast.name().name], ast) for ast in properties_asts]
+    return stubs
+
+
 def interface_script(mod_interface, nn_module):
     """
     Makes a ScriptModule from an nn.Module, using the interface methods rule for
@@ -612,7 +643,7 @@ def compile_unbound_method(concrete_type, fn):
     with torch._jit_internal._disable_emit_hooks():
         # We don't want to call the hooks here since the graph that is calling
         # this function is not yet complete
-        create_methods_from_stubs(concrete_type, (stub,))
+        create_methods_and_properties_from_stubs(concrete_type, (stub,), ())
     return stub
 
 def lazy_bind(concrete_type, unbound_method):
diff --git a/torch/jit/_script.py b/torch/jit/_script.py
diff --git a/torch/jit/frontend.py b/torch/jit/frontend.py
diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py
