PR comments

Author: zdevito

torch/csrc/jit/script/compiler.cpp

@@ -90,7 +90,9 @@ struct Environment {
 
     return sv;
   }
-
+  Block* block() {
+    return b;
+  }
   Symbol getBlockOwningKind() {
     Symbol owning_kind = Symbol();
     if (b->owningNode()) {
@@ -538,10 +540,18 @@ struct to_ir {
     // unrolled
     auto sv = emitSugaredExpr(itrs[0]);
     auto instances = sv->unrolledFor(stmt.range(), method);
+    const std::string& target_name = targets[0].name();
+    pushFrame(environment_stack->block());
     for(auto inst : instances) {
-      environment_stack->setSugaredVar(targets[0].name(), inst);
+      environment_stack->setSugaredVar(target_name, inst);
       emitStatements(body);
     }
+    for (const auto & n : environment_stack->definedVariables()) {
+      if (environment_stack->findInParentFrame(n)) {
+        environment_stack->next->setVar(n, environment_stack->getVar(n, stmt.range()));
+      }
+    }
+    popFrame();
   }
 
   void emitWhile(const While& stmt) {

torch/csrc/jit/script/compiler.h

@@ -46,7 +46,7 @@ struct SugaredValue : public std::enable_shared_from_this<SugaredValue> {
   }
 
   // use it in `for i in this: ...`
-  // in this case we will unroll the loop boy by assigning i to each of
+  // in this case we will unroll the loop body by assigning i to each of
   // the SugaredValues returned from this method.
   virtual std::vector<std::shared_ptr<SugaredValue>> unrolledFor(SourceRange loc, Method& m) {
     throw ErrorReport(loc) << kind() << " is not iterable";

torch/csrc/jit/script/init.cpp

@@ -88,10 +88,6 @@ struct VISIBILITY_HIDDEN PythonValue : public SugaredValue {
   py::object self;
 };
 
-// get the SugaredValue for something inside a torch.jit.Const
-// this can either be a ConstantPythonValue or a ModuleValue
-static std::shared_ptr<SugaredValue> createConstantSugaredValue(py::object obj);
-
 // by using torch.jit.Const, a user can mark a python value constant
 // we then make that value immutable.
 // once marked constant, we enable additional behavior such as
@@ -116,7 +112,7 @@ struct VISIBILITY_HIDDEN ConstantPythonValue : public PythonValue {
     py::tuple tup = self;
     std::vector<std::shared_ptr<SugaredValue>> result;
     for(size_t i = 0; i < tup.size(); ++i) {
-      result.push_back(createConstantSugaredValue(tup[i]));
+      result.push_back(std::make_shared<ConstantPythonValue>(tup[i]));
     }
     return result;
   }
@@ -195,7 +191,7 @@ struct ModuleValue : public SugaredValue {
          py::isinstance(attr, py::module::import("torch.nn").attr("Module"))) {
         return std::make_shared<PythonValue>(attr);
       } else if(py_module.attr("_constants_set").contains(field.c_str())) {
-        return createConstantSugaredValue(attr);
+        return std::make_shared<ConstantPythonValue>(attr);
       } else {
         throw ErrorReport(loc) << "attribute '" << field << "' of type '" << typeString(attr) << "' is not usable in a script method (did you forget to add it __constants__?)";
       }
@@ -213,8 +209,13 @@ struct ModuleValue : public SugaredValue {
       return SugaredValue::unrolledFor(loc, m);
     std::vector<std::shared_ptr<SugaredValue>> result;
     for(py::handle module : py_module) {
-      result.push_back(createConstantSugaredValue(
-        py::reinterpret_borrow<py::object>(module)));
+      py::object obj = py::reinterpret_borrow<py::object>(module);
+      if(py::isinstance<Module>(obj)) {
+        auto r = py::cast<std::shared_ptr<Module>>(obj);
+        result.push_back(std::make_shared<ModuleValue>(r));
+      } else {
+        result.push_back(std::make_shared<ConstantPythonValue>(obj));
+      }
     }
     return result;
   }
@@ -223,13 +224,6 @@ struct ModuleValue : public SugaredValue {
   std::shared_ptr<Module> module;
 };
 
-static std::shared_ptr<SugaredValue> createConstantSugaredValue(py::object obj) {
-  if(py::isinstance<Module>(obj)) {
-    auto r = py::cast<std::shared_ptr<Module>>(obj);
-    return std::make_shared<ModuleValue>(r);
-  }
-  return std::make_shared<ConstantPythonValue>(obj);
-}
 
 // TODO: dedup with other init
 

torch/jit/__init__.py

@@ -656,14 +656,15 @@ def __getitem__(self, k):
             raise KeyError(k)
         return self.module._get_parameter(k)
 
-# base types that can be constants in addition to tuples and lists
+# 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
 # ConstantValue in jit/script/init.cpp
-_constant_types = [bool, float, int, types.FunctionType]
+_constant_types = (bool, float, int, types.FunctionType)
 
 
 def _get_valid_constant(v):
-    if any(isinstance(v, typ) for typ in _constant_types):
+    if isinstance(v, _constant_types):
         return v
     elif isinstance(v, tuple) or isinstance(v, list):
         return tuple(_get_valid_constant(x) for x in v)