feat: Add converter for aten::isfinite (#1841)

Author: mfeliz-cruise

core/conversion/converters/impl/unary.cpp

@@ -79,6 +79,26 @@ auto logical_not_registration TORCHTRT_UNUSED = RegisterNodeConversionPatterns()
        return true;
      }});
 
+auto isfinite_registration TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
+    {"aten::isfinite(Tensor self) -> Tensor", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+       auto in = args[0].ITensorOrFreeze(ctx);
+       // assuming x-x = 0 for all values other than nan/inf/-inf where x-x = nan
+       // x==x for all non-nan values
+       auto inf_test_layer = ctx->net->addElementWise(*in, *in, nvinfer1::ElementWiseOperation::kSUB);
+       TORCHTRT_CHECK(inf_test_layer, "Unable to create sub layer from node: " << *n);
+       inf_test_layer->setName((util::node_info(n) + "_inf_test").c_str());
+       auto inf_test_tensor = inf_test_layer->getOutput(0);
+
+       auto nan_test_layer =
+           ctx->net->addElementWise(*inf_test_tensor, *inf_test_tensor, nvinfer1::ElementWiseOperation::kEQUAL);
+       TORCHTRT_CHECK(nan_test_layer, "Unable to create eq layer from node: " << *n);
+       nan_test_layer->setName((util::node_info(n) + "_nan_test").c_str());
+
+       auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], nan_test_layer->getOutput(0));
+       LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+       return true;
+     }});
+
 #define convert(unary, trt_type)                                                               \
   auto unary##_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(       \
       {"aten::" #unary "(Tensor self) -> Tensor",                                              \

tests/core/conversion/converters/test_unary.cpp

@@ -111,6 +111,27 @@ TEST(Converters, ATenLogicalNotBoolConvertsCorrectly) {
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
 }
 
+TEST(Converters, ATenFiniteConvertsCorrectly) {
+  const auto graph = gen_test_graph("isfinite");
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+  auto in = at::tensor(
+      {float(0),
+       std::nanf(""),
+       float(2),
+       std::numeric_limits<float>::infinity(),
+       float(4),
+       -std::numeric_limits<float>::infinity()},
+      {at::kCUDA});
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in});
+
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in});
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
+}
+
 #define test_unary(unary, name)                                                                  \
   TEST(Converters, ATen##name##ConvertsCorrectly) {                                              \
     const auto graph = gen_test_graph(#unary);                                                   \