Make changes to scale implementation

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

core/conversion/converters/impl/element_wise.cpp

@@ -153,13 +153,14 @@ auto element_wise_registrations TRTORCH_UNUSED =
                     auto other = args[1].ITensorOrFreeze(ctx);
 
                     if (1 != scalar) {
-                      auto scaleW = Weights(ctx, scalar);
-                      auto unuse = Weights();
-                      // IScaleLayer assert shift, scale and power to have
-                      // the same dtype
-                      auto scaleLayer = ctx->net->addScale(
-                          *other, nvinfer1::ScaleMode::kUNIFORM, unuse.data, scaleW.data, unuse.data);
-                      TRTORCH_CHECK(scaleLayer, "Unable to create scale layer from node: " << *n);
+                      auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+                      auto scaleLayer = add_elementwise(
+                          ctx,
+                          nvinfer1::ElementWiseOperation::kPROD,
+                          other,
+                          alphaTensor,
+                          util::node_info(n) + std::string("_AlphaMultiplier"));
+                      TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
                       other = scaleLayer->getOutput(0);
                     }
 
@@ -181,13 +182,14 @@ auto element_wise_registrations TRTORCH_UNUSED =
                     auto other = args[1].ITensorOrFreeze(ctx);
 
                     if (1 != scalar) {
-                      auto scaleW = Weights(ctx, scalar);
-                      auto unuse = Weights();
-                      // IScaleLayer assert shift, scale and power to have
-                      // the same dtype
-                      auto scaleLayer = ctx->net->addScale(
-                          *other, nvinfer1::ScaleMode::kUNIFORM, unuse.data, scaleW.data, unuse.data);
-                      TRTORCH_CHECK(scaleLayer, "Unable to create scale layer from node: " << *n);
+                      auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+                      auto scaleLayer = add_elementwise(
+                          ctx,
+                          nvinfer1::ElementWiseOperation::kPROD,
+                          other,
+                          alphaTensor,
+                          util::node_info(n) + std::string("_AlphaMultiplier"));
+                      TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
                       other = scaleLayer->getOutput(0);
                     }
 
@@ -209,13 +211,14 @@ auto element_wise_registrations TRTORCH_UNUSED =
                     auto scalar = args[2].unwrapToScalar().to<float>();
 
                     if (1 != scalar) {
-                      auto scaleW = Weights(ctx, scalar);
-                      auto unuse = Weights();
-                      // IScaleLayer assert shift, scale and power to have
-                      // the same dtype
-                      auto scaleLayer =
-                          ctx->net->addScale(*self, nvinfer1::ScaleMode::kUNIFORM, unuse.data, scaleW.data, unuse.data);
-                      TRTORCH_CHECK(scaleLayer, "Unable to create scale layer from node: " << *n);
+                      auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+                      auto scaleLayer = add_elementwise(
+                          ctx,
+                          nvinfer1::ElementWiseOperation::kPROD,
+                          self,
+                          alphaTensor,
+                          util::node_info(n) + std::string("_AlphaMultiplier"));
+                      TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
                       self = scaleLayer->getOutput(0);
                     }
 
@@ -236,13 +239,14 @@ auto element_wise_registrations TRTORCH_UNUSED =
                     auto scalar = args[2].unwrapToScalar().to<float>();
 
                     if (1 != scalar) {
-                      auto scaleW = Weights(ctx, scalar);
-                      auto unuse = Weights();
-                      // IScaleLayer assert shift, scale and power to have
-                      // the same dtype
-                      auto scaleLayer =
-                          ctx->net->addScale(*self, nvinfer1::ScaleMode::kUNIFORM, unuse.data, scaleW.data, unuse.data);
-                      TRTORCH_CHECK(scaleLayer, "Unable to create scale layer from node: " << *n);
+                      auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+                      auto scaleLayer = add_elementwise(
+                          ctx,
+                          nvinfer1::ElementWiseOperation::kPROD,
+                          self,
+                          alphaTensor,
+                          util::node_info(n) + std::string("_AlphaMultiplier"));
+                      TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
                       self = scaleLayer->getOutput(0);
                     }
 

tests/core/conversion/converters/test_element_wise.cpp

@@ -89,7 +89,7 @@ TEST(Converters, ATenAddWithScalarConvertsCorrectly) {
 TEST(Converters, ATenSubConvertsCorrectly) {
   const auto graph = R"IR(
       graph(%0 : Tensor, %1 : Tensor):
-        %2 : int = prim::Constant[value=1]()
+        %2 : int = prim::Constant[value=2.3]()
         %3 : Tensor = aten::sub(%0, %1, %2)
         return (%3))IR";
   pointwise_test_helper(graph, false);
@@ -215,6 +215,8 @@ TEST(Converters, ATenRsubWithTensorConvertsCorrectly) {
         %2 : int = prim::Constant[value=2]()
         %3 : Tensor = aten::rsub(%0, %1, %2)
         return (%3))IR";
+  pointwise_test_helper(graph, false, false, {3, 4}, {4});
+  pointwise_test_helper(graph, false, false, {4}, {3, 4});
   pointwise_test_helper(graph, false, true, {4, 3, 3, 3}, {4, 3, 3, 3});
 }
 
@@ -226,4 +228,4 @@ TEST(Converters, ATenRsubWithScalarConvertsCorrectly) {
         %3 : Tensor = aten::rsub(%0, %scalar, %2)
         return (%3))IR";
   pointwise_test_helper(graph, true, false, {4, 3, 3, 3});
-}
+}