feat(): started to work on add_.t evaluator, doesn't work yet

Signed-off-by: Abhiram Iyer <[email protected]> Signed-off-by: Abhiram Iyer <[email protected]>
pytorch · Jun 29, 2020 · f216d3f · f216d3f
1 parent f594e43
commit f216d3f
Showing 1 changed file with 24 additions and 0 deletions.
diff --git a/core/conversion/evaluators/aten.cpp b/core/conversion/evaluators/aten.cpp
@@ -10,6 +10,8 @@
 #include "core/conversion/evaluators/evaluators.h"
 #include "core/conversion/evaluators/eval_macros.h"
 
+// #include <csignal>
+
 namespace trtorch {
 namespace core {
 namespace conversion {
@@ -243,6 +245,28 @@ auto aten_registrations TRTORCH_UNUSED = RegisterNodeEvaluators()
             "aten::add.int(int a, int b) -> (int)",
             "aten::add.float(float a, float b) -> (float)"
         })
+    }).evaluator({
+        c10::Symbol::fromQualString("aten::add_"),
+        [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+            LOG_DEBUG("aten::add_ evaluator is found");
+
+            // std::raise(SIGINT);
+
+            if (args.at(n->input(0)).IValue()->isList()) {
+                auto a = args.at(n->input(0)).IValue()->to<c10::List<c10::IValue>>();
+                auto b = args.at(n->input(1)).IValue()->to<c10::List<c10::IValue>>();
+
+                // incorrect syntax
+                // for (auto each : b) {
+                //     a.push_back(each);
+                // }
+
+                return a;
+            }
+        },
+        EvalOptions().validSchemas({
+            "aten::add_.t(t[](a!) self, t[] b) -> (t[])"
+        })
     }).evaluator({
         c10::Symbol::fromQualString("aten::mul"),
         [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {