Use a stable lowering for ReduceLogSumExp

onnxOpImporters.cpp

@@ -5011,16 +5011,61 @@ DEFINE_BUILTIN_OP_IMPORTER(ReduceLogSumExp)
         RETURN_IDENTITY(inputs.at(0), node, nodeIdx);
     }
 
-    std::vector<TensorOrWeights> expResult
-        = unaryHelper(ctx, node, nodeIdx, inputs.at(0), nvinfer1::UnaryOperation::kEXP);
+    TensorOrWeights input = inputs.at(0);
+    TensorOrWeights inputAxes = inputs.size() >= 2 ? inputs.at(1) : TensorOrWeights();
 
-    // Include the axes input if present to ensure reduction is performed correctly
-    if (inputs.size() >= 2)
+    OnnxAttrs attrs(node, ctx);
+    bool const keepdims = attrs.get("keepdims", 1);
+    int32_t const ndim = input.shape().nbDims;
+
+    std::vector<int32_t> axes;
+    if (attrs.count("axes"))
+    {
+        axes = attrs.get<std::vector<int32_t>>("axes");
+    }
+    else if (!inputAxes.isNullTensor())
+    {
+        ONNXTRT_CHECK_NODE(
+            inputAxes.is_weights(), "Axis input must be an initializer!", node, nodeIdx, ErrorCode::kUNSUPPORTED_NODE);
+        weightsToVector<int32_t>(inputAxes.weights(), &axes);
+    }
+    if (axes.empty())
+    {
+        if (attrs.get("noop_with_empty_axes", 0) == 1)
+        {
+            RETURN_IDENTITY(inputs.at(0), node, nodeIdx);
+        }
+        axes.resize(ndim);
+        std::iota(axes.begin(), axes.end(), 0);
+    }
+    for (int32_t& axis : axes)
     {
-        expResult.push_back(inputs.at(1));
+        convertAxis(axis, ndim, node, nodeIdx);
     }
+    std::sort(axes.begin(), axes.end());
 
-    return importReduceLogSum(ctx, node, nodeIdx, expResult);
+    auto maxResult = reduceTensor(ctx, node, nodeIdx, input, nvinfer1::ReduceOperation::kMAX, inputAxes);
+    nvinfer1::ITensor* reducedMax = &convertToTensor(maxResult.at(0), ctx);
+
+    nvinfer1::ITensor* maxForSub = reducedMax;
+    if (!keepdims)
+    {
+        maxForSub = unsqueezeTensor(ctx, *reducedMax, axes);
+    }
+
+    auto& inputTensor = convertToTensor(input, ctx);
+    auto* shifted = getElementWiseResult(ctx, inputTensor, *maxForSub, nvinfer1::ElementWiseOperation::kSUB);
+    auto* shiftedExp = getUnaryResult(ctx, *shifted, nvinfer1::UnaryOperation::kEXP);
+
+    std::vector<TensorOrWeights> sumInputs{TensorOrWeights{shiftedExp}};
+    if (inputs.size() >= 2)
+    {
+        sumInputs.push_back(inputAxes);
+    }
+    auto sumResult = importReduceSum(ctx, node, nodeIdx, sumInputs);
+    auto* logSum = getUnaryResult(ctx, convertToTensor(sumResult.at(0), ctx), nvinfer1::UnaryOperation::kLOG);
+    auto* stableResult = getElementWiseResult(ctx, *logSum, *reducedMax, nvinfer1::ElementWiseOperation::kSUM);
+    return {{stableResult}};
 }
 DECLARE_BUILTIN_OP_IMPORTER(ReduceSumSquare);
 DEFINE_BUILTIN_OP_IMPORTER(ReduceL2)

onnx_backend_test.py

@@ -9,6 +9,9 @@
 
 import unittest
 import onnx.backend.test
+import numpy as np
+from onnx import TensorProto
+from onnx import helper as onnx_helper
 
 import onnx_tensorrt.backend as trt
 
@@ -176,6 +179,69 @@
 # dilations not supported in ConvTRanspose layer
 backend_test.exclude(r'.*test_convtranspose_dilations_custom_cuda')
 
+
+class TensorRTCustomReduceLogSumExpTest(unittest.TestCase):
+    def test_reduce_log_sum_exp_large_finite_input_custom(self):
+        node = onnx_helper.make_node("ReduceLogSumExp", ["x"], ["y"], keepdims=1)
+        graph = onnx_helper.make_graph(
+            [node],
+            "reduce_log_sum_exp_large_finite_input_custom",
+            [onnx_helper.make_tensor_value_info("x", TensorProto.FLOAT, [4])],
+            [onnx_helper.make_tensor_value_info("y", TensorProto.FLOAT, [1])],
+        )
+        model = onnx_helper.make_model(graph, opset_imports=[onnx_helper.make_opsetid("", 18)])
+        model.ir_version = 10
+
+        x = np.array([250.0, 248.0, 255.0, 251.0], dtype=np.float32)
+        expected = np.log(np.sum(np.exp(x - np.max(x)))) + np.max(x)
+
+        outputs = trt.run_model(model, [x], device="CUDA:0")
+        actual = float(outputs.y[0])
+
+        self.assertTrue(np.isfinite(actual))
+        self.assertAlmostEqual(actual, float(expected), places=5)
+
+    def test_reduce_log_sum_exp_keepdims_zero_custom(self):
+        node = onnx_helper.make_node("ReduceLogSumExp", ["x", "axes"], ["y"], keepdims=0)
+        graph = onnx_helper.make_graph(
+            [node],
+            "reduce_log_sum_exp_keepdims_zero_custom",
+            [onnx_helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 4])],
+            [onnx_helper.make_tensor_value_info("y", TensorProto.FLOAT, [2])],
+            initializer=[onnx_helper.make_tensor("axes", TensorProto.INT64, [1], np.array([1], dtype=np.int64))],
+        )
+        model = onnx_helper.make_model(graph, opset_imports=[onnx_helper.make_opsetid("", 18)])
+        model.ir_version = 10
+
+        x = np.array([[250.0, 248.0, 255.0, 251.0], [100.0, 101.0, 99.0, 97.0]], dtype=np.float32)
+        expected = np.log(np.sum(np.exp(x - np.max(x, axis=1, keepdims=True)), axis=1)) + np.max(x, axis=1)
+
+        outputs = trt.run_model(model, [x], device="CUDA:0")
+        actual = np.asarray(outputs.y)
+
+        self.assertTrue(np.all(np.isfinite(actual)))
+        np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5)
+
+    def test_reduce_log_sum_exp_noop_empty_axes_custom(self):
+        node = onnx_helper.make_node("ReduceLogSumExp", ["x", "axes"], ["y"], keepdims=0, noop_with_empty_axes=1)
+        graph = onnx_helper.make_graph(
+            [node],
+            "reduce_log_sum_exp_noop_empty_axes_custom",
+            [onnx_helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 2])],
+            [onnx_helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 2])],
+            initializer=[onnx_helper.make_tensor("axes", TensorProto.INT64, [0], np.array([], dtype=np.int64))],
+        )
+        model = onnx_helper.make_model(graph, opset_imports=[onnx_helper.make_opsetid("", 18)])
+        model.ir_version = 10
+
+        x = np.array([[250.0, 248.0], [255.0, 251.0]], dtype=np.float32)
+
+        outputs = trt.run_model(model, [x], device="CUDA:0")
+        actual = np.asarray(outputs.y)
+
+        self.assertEqual(actual.size, x.size)
+        np.testing.assert_allclose(actual.reshape(x.shape), x, rtol=1e-6, atol=1e-6)
+
 globals().update(backend_test
                  .enable_report()
                  .test_cases)