add test case for aten::div evaluator

Signed-off-by: inocsin <[email protected]>

Author: inocsin

core/conversion/evaluators/aten.cpp

@@ -427,7 +427,7 @@ auto aten_registrations TRTORCH_UNUSED =
                     },
                     EvalOptions().validSchemas({
                         "aten::div.float(float a, float b) -> (float)",
-                        "aten::div.int(int a, int b) -> (int)",
+                        "aten::div.int(int a, int b) -> (float)",
                     })})
         .evaluator({c10::Symbol::fromQualString("aten::floordiv"),
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {

tests/core/conversion/evaluators/BUILD

@@ -11,9 +11,14 @@ evaluator_test(
   name = "test_prim_evaluators",
 )
 
+evaluator_test(
+  name = "test_aten_evaluators",
+)
+
 test_suite(
     name = "evaluator_tests",
     tests = [
-        ":test_prim_evaluators"
+        ":test_prim_evaluators",
+        ":test_aten_evaluators"
     ]
 )

tests/core/conversion/evaluators/test_aten_evaluators.cpp

@@ -0,0 +1,39 @@
+#include <string>
+#include "core/compiler.h"
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "torch/csrc/jit/ir/irparser.h"
+
+TEST(Evaluators, DivIntEvaluatesCorrectly) {
+  const auto graph = R"IR(
+      graph():
+        %1 : int = prim::Constant[value=9]()
+        %2 : int = prim::Constant[value=4]()
+        %3 : float = aten::div(%1, %2)
+        return (%3))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto jit_results = trtorch::tests::util::EvaluateGraphJIT(g, {});
+  auto trt_results = trtorch::tests::util::EvaluateGraph(g->block(), {});
+
+  ASSERT_TRUE(jit_results[0] == trt_results[0]);
+}
+
+TEST(Evaluators, DivFloatEvaluatesCorrectly) {
+  const auto graph = R"IR(
+      graph():
+        %1 : float = prim::Constant[value=9.1]()
+        %2 : float = prim::Constant[value=4.2]()
+        %3 : float = aten::div(%1, %2)
+        return (%3))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto jit_results = trtorch::tests::util::EvaluateGraphJIT(g, {});
+  auto trt_results = trtorch::tests::util::EvaluateGraph(g->block(), {});
+
+  ASSERT_TRUE(jit_results[0] == trt_results[0]);
+}