Added test cases for aten::ScalarImplicit

- Resolved issue with Tensor casting issue
- Added tests to cover int, float, and complex cases

Author: gs-olive

tests/core/lowering/test_device_casting.cpp

@@ -5,6 +5,7 @@
 #include "gtest/gtest.h"
 #include "tests/util/util.h"
 #include "torch/csrc/jit/ir/irparser.h"
+#include "torch/csrc/jit/ir/subgraph_matcher.h"
 #include "torch/csrc/jit/passes/common_subexpression_elimination.h"
 #include "torch/torch.h"
 
@@ -36,8 +37,8 @@ TEST(LoweringPasses, UnpackAndCastNumToTensorLowersIntCorrectly) {
         %2 : Tensor = prim::NumToTensor(%x.1)
         return (%2))IR";
 
-  // Make range [0.01, 1.01] to ensure positives / avoid NaN with negative sqrt
   auto in = 1;
+
   auto g = std::make_shared<torch::jit::Graph>();
   torch::jit::parseIR(graph, g.get());
 
@@ -116,3 +117,78 @@ TEST(LoweringPasses, UnpackAndCastFullFloatLowersCorrectly) {
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(
       jit_pre_results[0].toTensor(), jit_post_results[0].toTensor().cpu(), 2e-6));
 }
+
+TEST(LoweringPasses, ReplaceScalarImplicitLowersCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1: Tensor):
+        %5 : int = prim::Constant[value=0]()
+        %false : bool = prim::Constant[value=0]()
+        %none : NoneType = prim::Constant()
+        %cuda : Device = prim::Constant[value="cuda"]()
+        %3 : int = aten::size(%x.1, %5)
+        %y.2 : Tensor = prim::NumToTensor(%3)
+        %y.1 : Tensor = aten::to(%y.2, %cuda, %none, %false, %false)
+        %19 : Tensor[] = prim::ListConstruct(%x.1, %y.1)
+        %21 : Tensor, %22 : Tensor = prim::ListUnpack(%19)
+        %2 : Scalar = aten::ScalarImplicit(%22)
+        %out : Tensor = prim::NumToTensor(%2)
+        return (%out))IR";
+
+  auto in = at::rand({2, 3, 5, 7}, {at::kCUDA});
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto jit_pre_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+  torch_tensorrt::core::lowering::passes::ReplaceScalarImplicit(g);
+  torch::jit::EliminateCommonSubexpression(g);
+  auto jit_post_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_pre_results[0].toTensor(), jit_post_results[0].toTensor(), 2e-6));
+}
+
+TEST(LoweringPasses, ReplaceScalarImplicitIntNumToTensorLowersCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1: int):
+        %1 : Tensor = prim::NumToTensor(%x.1)
+        %2 : Scalar = aten::ScalarImplicit(%1)
+        %3 : Tensor = prim::NumToTensor(%2)
+        return (%3))IR";
+
+  auto in = 25;
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto jit_pre_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+  torch_tensorrt::core::lowering::passes::UnpackAndCastNumToTensor(g);
+  torch_tensorrt::core::lowering::passes::ReplaceScalarImplicit(g);
+  torch::jit::EliminateCommonSubexpression(g);
+  auto jit_post_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_pre_results[0].toTensor(), jit_post_results[0].toTensor(), 2e-6));
+}
+
+TEST(LoweringPasses, ReplaceScalarImplicitFloatLowersCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1: float):
+        %1 : Tensor = prim::NumToTensor(%x.1)
+        %2 : Scalar = aten::ScalarImplicit(%1)
+        %3 : Tensor = prim::NumToTensor(%2)
+        return (%3))IR";
+
+  auto in = 2.5;
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto jit_pre_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+  torch_tensorrt::core::lowering::passes::ReplaceScalarImplicit(g);
+  torch::jit::EliminateCommonSubexpression(g);
+  auto jit_post_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {in});
+
+  ASSERT_TRUE(
+      torch_tensorrt::tests::util::almostEqual(jit_pre_results[0].toTensor(), jit_post_results[0].toTensor(), 2e-6));
+}