Merge pull request #300 from NVIDIA/ne

Add support for ne layer

Author: narendasan

core/conversion/converters/impl/element_wise.cpp

@@ -260,6 +260,70 @@ auto element_wise_registrations TRTORCH_UNUSED =
                     LOG_DEBUG("Output tensor shape: " << out->getDimensions());
                     return true;
                   }})
+        .pattern({"aten::ne.Tensor(Tensor self, Tensor other) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    // TODO: Remove with functionalization
+                    auto self = args[0].ITensorOrFreeze(ctx);
+                    auto other = args[1].ITensorOrFreeze(ctx);
+                    auto equal = add_elementwise(
+                        ctx,
+                        nvinfer1::ElementWiseOperation::kEQUAL,
+                        self,
+                        other,
+                        util::node_info(n) + std::string("is_equal"));
+                    TRTORCH_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n);
+                    // XOR with ones negates and produces not_equal result
+                    auto options = torch::TensorOptions().dtype(torch::kFloat32);
+                    auto ones = at::full({1}, 1, {options});
+                    auto ones_tensor = tensor_to_const(ctx, ones);
+                    nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor);
+                    cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL);
+
+                    auto sub = add_elementwise(
+                        ctx,
+                        nvinfer1::ElementWiseOperation::kXOR,
+                        cast_layer->getOutput(0),
+                        equal->getOutput(0),
+                        util::node_info(n));
+                    TRTORCH_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n);
+
+                    sub->setName(util::node_info(n).c_str());
+                    auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0));
+                    LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions());
+                    return true;
+                  }})
+        .pattern({"aten::ne.Scalar(Tensor self, Scalar other) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    auto self = args[0].ITensorOrFreeze(ctx);
+                    auto scalar = args[1].unwrapToScalar().to<float>();
+                    auto scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar}));
+                    auto equal = add_elementwise(
+                        ctx,
+                        nvinfer1::ElementWiseOperation::kEQUAL,
+                        self,
+                        scalar_tensor,
+                        util::node_info(n) + std::string("is_equal"));
+                    TRTORCH_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n);
+                    // XOR with ones negates and produces not_equal result
+                    auto options = torch::TensorOptions().dtype(torch::kFloat32);
+                    auto ones = at::full({1}, 1, {options});
+                    auto ones_tensor = tensor_to_const(ctx, ones);
+                    nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor);
+                    cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL);
+
+                    auto sub = add_elementwise(
+                        ctx,
+                        nvinfer1::ElementWiseOperation::kXOR,
+                        cast_layer->getOutput(0),
+                        equal->getOutput(0),
+                        util::node_info(n));
+                    TRTORCH_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n);
+
+                    sub->setName(util::node_info(n).c_str());
+                    auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0));
+                    LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions());
+                    return true;
+                  }})
         .pattern({"aten::pow.Tensor_Tensor(Tensor self, Tensor exponent) -> (Tensor)",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     // TODO: Remove with functionalization

tests/core/conversion/converters/test_element_wise.cpp

@@ -76,6 +76,16 @@ TEST(Converters, ATenAddImplicitWithAlphaConvertsCorrectly) {
   pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3});
 }
 
+TEST(Converters, ATenAddWithScalarConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=1]()
+        %scalar : float = prim::Constant[value=2.4]()
+        %3 : Tensor = aten::add(%0, %scalar, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, true);
+}
+
 TEST(Converters, ATenSubConvertsCorrectly) {
   const auto graph = R"IR(
       graph(%0 : Tensor, %1 : Tensor):
@@ -134,12 +144,22 @@ TEST(Converters, ATenPowScalarConvertsCorrectly) {
   pointwise_test_helper(graph, true);
 }
 
-TEST(Converters, ATenAddWithScalarConvertsCorrectly) {
+TEST(Converters, ATenNeTensorConvertsCorrectly) {
   const auto graph = R"IR(
-      graph(%0 : Tensor):
-        %2 : int = prim::Constant[value=1]()
-        %scalar : float = prim::Constant[value=2.4]()
-        %3 : Tensor = aten::add(%0, %scalar, %2)
+    graph(%x.1 : Tensor,
+      %y.1 : Tensor):
+        %3 : Tensor = aten::ne(%x.1, %y.1)
         return (%3))IR";
-  pointwise_test_helper(graph, true);
+  pointwise_test_helper(graph, false, false, {3, 4}, {3, 4});
+  pointwise_test_helper(graph, false, true, {3, 4}, {3, 4});
+}
+
+TEST(Converters, ATenNeScalarConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%x.1 : Tensor):
+            %2 : int = prim::Constant[value=2]()
+            %3 : Tensor = aten::ne(%x.1, %2)
+            return (%3))IR";
+  pointwise_test_helper(graph, true, false, {3, 4, 2});
+  ;
 }