Improve error message in IR parser when accessing undefined variable.

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/26771

Test Plan: Imported from OSS

Differential Revision: D17562853

Pulled By: ZolotukhinM

fbshipit-source-id: b4d4bc6001e3ea06f4d1b8691ad2a339a04c16ea
diff --git a/torch/csrc/jit/irparser.cpp b/torch/csrc/jit/irparser.cpp
index a1a2f04..5b91fdf 100644
--- a/torch/csrc/jit/irparser.cpp
+++ b/torch/csrc/jit/irparser.cpp
@@ -55,6 +55,8 @@
       int end,
       const std::function<void()>& callback);
 
+  Value* findValueInVMap(const std::string& name);
+
   torch::jit::script::Lexer L;
   torch::jit::Graph* g = nullptr;
   std::unordered_map<std::string, Value*>& vmap;
@@ -255,8 +257,7 @@
   }
   parseList('(', ',', ')', [&] {
     std::string var_name = parseVar();
-    AT_ASSERT(vmap.count(var_name));
-    n->addInput(vmap[var_name]);
+    n->addInput(findValueInVMap(var_name));
   });
 }
 
@@ -282,8 +283,7 @@
   L.expect(TK_ARROW);
   parseList('(', ',', ')', [&] {
     std::string var_name = parseVar();
-    AT_ASSERT(vmap.count(var_name));
-    b->registerOutput(vmap[var_name]);
+    b->registerOutput(findValueInVMap(var_name));
   });
   L.expect(TK_NEWLINE);
   L.expect(TK_DEDENT);
@@ -385,10 +385,7 @@
   // Parse output names and types
   parseList('(', ',', ')', [&] {
     std::string var_name = parseVar();
-    // Outputs should already be in VMAP, otherwise we're trying to return
-    // undefined value.
-    AT_ASSERT(vmap.count(var_name));
-    g->registerOutput(vmap.at(var_name));
+    g->registerOutput(findValueInVMap(var_name));
   });
 
   // Consume ending tokens
@@ -439,6 +436,15 @@
     L.expect(end);
   }
 }
+
+Value* IRParser::findValueInVMap(const std::string& name) {
+  if (!vmap.count(name)) {
+    throw ErrorReport(L.cur().range)
+        << "Cannot find a variable with name '" << name << "'";
+  }
+  return vmap.at(name);
+}
+
 } // namespace script
 } // namespace jit
 } // namespace torch