[JIT] Add property support for ScriptModules

Meghan Lele · Meghan Lele · commit a63397a2ec1f · 2020-08-05T18:04:00.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: 4f44e0e Pull Request resolved: #42390
diff --git a/test/jit/test_recursive_script.py b/test/jit/test_recursive_script.py
@@ -607,24 +607,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,30 @@ 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
+
+            @attr.setter
+            def attr(self, a: int):
+                if a > 0:
+                    self.a = a
+                else:
+                    self.a = 0
+
+        self.checkModule(ModuleWithProperties(5), (5, 6,))
+        self.checkModule(ModuleWithProperties(5), (-5, -6,))
+
     def test_export_opnames_interface(self):
         global OneTwoModule
 
diff --git a/torch/csrc/jit/python/python_sugared_value.cpp b/torch/csrc/jit/python/python_sugared_value.cpp
@@ -589,6 +589,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
@@ -162,7 +162,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 property.setter().get().name();
+        }
+
+        return Ident::create(property.range(), "");
+      });
 
   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
@@ -1559,27 +1559,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, defs, resolvers, &self);
+            cu->define(
+                prefix,
+                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 = [
@@ -29,6 +31,16 @@
     "dump_patches",
 ]
 
+ignored_properties = [
+    # Temporary fix for RNN module property named 'all_weights' being scripted
+    "all_weights",
+    "original_name",
+    "graph",
+    "inlined_graph",
+    "code",
+    "code_with_constants",
+]
+
 def make_stub(func, name):
     rcb = _jit_internal.createResolutionCallbackFromClosure(func)
     ast = get_jit_def(func, name, self_name="RecursiveScriptModule")
@@ -48,6 +60,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 +252,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 +290,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 +337,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:
@@ -373,7 +385,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 +403,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 +419,20 @@ 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
 
 
+    for property_stub in property_stubs:
+        property_name = property_stub.def_.name().name
+        fget = cpp_module._get_method(property_stub.def_.getter_name().name)
+        fset = cpp_module._get_method(property_stub.def_.setter_name().name)
+        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 +566,32 @@ 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) and name not in ignored_properties:
+            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 = []
+    for ast in properties_asts:
+        if ast.name().name not in ignored_properties:
+            stubs.append(PropertyStub(rcbs[ast.name().name], ast))
+
+    return stubs
+
+
 def interface_script(mod_interface, nn_module):
     """
     Makes a ScriptModule from an nn.Module, using the interface methods rule for
@@ -612,7 +656,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):
