From 4249e0379abdbb436c7cd022806293b601f93dcc Mon Sep 17 00:00:00 2001 From: peter941221 Date: Wed, 10 Jun 2026 12:05:07 +0800 Subject: [PATCH 1/4] Stabilize ReduceLogSumExp lowering Signed-off-by: peter941221 --- onnxOpImporters.cpp | 53 +++++++++++++++++++++++++++++++++++++++----- onnx_backend_test.py | 25 +++++++++++++++++++++ 2 files changed, 72 insertions(+), 6 deletions(-) diff --git a/onnxOpImporters.cpp b/onnxOpImporters.cpp index 521be92e..978a6b6a 100644 --- a/onnxOpImporters.cpp +++ b/onnxOpImporters.cpp @@ -5011,16 +5011,57 @@ DEFINE_BUILTIN_OP_IMPORTER(ReduceLogSumExp) RETURN_IDENTITY(inputs.at(0), node, nodeIdx); } - std::vector 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) + auto maxResult = reduceTensor(ctx, node, nodeIdx, input, nvinfer1::ReduceOperation::kMAX, inputAxes); + nvinfer1::ITensor* reducedMax = &convertToTensor(maxResult.at(0), ctx); + + OnnxAttrs attrs(node, ctx); + bool const keepdims = attrs.get("keepdims", 1); + int32_t const ndim = input.shape().nbDims; + + std::vector axes; + if (attrs.count("axes")) + { + axes = attrs.get>("axes"); + } + else if (!inputAxes.isNullTensor()) + { + ONNXTRT_CHECK_NODE( + inputAxes.is_weights(), "Axis input must be an initializer!", node, nodeIdx, ErrorCode::kUNSUPPORTED_NODE); + weightsToVector(inputAxes.weights(), &axes); + } + if (axes.empty()) { - expResult.push_back(inputs.at(1)); + axes.resize(ndim); + std::iota(axes.begin(), axes.end(), 0); + } + for (int32_t& axis : axes) + { + convertAxis(axis, ndim, node, nodeIdx); + } + std::sort(axes.begin(), axes.end()); + + nvinfer1::ITensor* maxForSub = reducedMax; + if (!keepdims) + { + maxForSub = unsqueezeTensor(ctx, *reducedMax, axes); } - return importReduceLogSum(ctx, node, nodeIdx, expResult); + auto& inputTensor = convertToTensor(input, ctx); + auto* shifted = getElementWiseResult(ctx, inputTensor, *maxForSub, nvinfer1::ElementWiseOperation::kSUB); + auto* shiftedExp = getUnaryResult(ctx, *shifted, nvinfer1::UnaryOperation::kEXP); + + std::vector 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) diff --git a/onnx_backend_test.py b/onnx_backend_test.py index f62004ad..6616452b 100644 --- a/onnx_backend_test.py +++ b/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,28 @@ # 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) + globals().update(backend_test .enable_report() .test_cases) From 5d9f6ad74005ec019466b4a65fedac998830b14b Mon Sep 17 00:00:00 2001 From: peter941221 Date: Wed, 10 Jun 2026 12:27:58 +0800 Subject: [PATCH 2/4] Handle noop empty axes in ReduceLogSumExp Signed-off-by: peter941221 --- onnxOpImporters.cpp | 10 +++++++--- onnx_backend_test.py | 40 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 3 deletions(-) diff --git a/onnxOpImporters.cpp b/onnxOpImporters.cpp index 978a6b6a..b3520a8f 100644 --- a/onnxOpImporters.cpp +++ b/onnxOpImporters.cpp @@ -5014,9 +5014,6 @@ DEFINE_BUILTIN_OP_IMPORTER(ReduceLogSumExp) TensorOrWeights input = inputs.at(0); TensorOrWeights inputAxes = inputs.size() >= 2 ? inputs.at(1) : TensorOrWeights(); - auto maxResult = reduceTensor(ctx, node, nodeIdx, input, nvinfer1::ReduceOperation::kMAX, inputAxes); - nvinfer1::ITensor* reducedMax = &convertToTensor(maxResult.at(0), ctx); - OnnxAttrs attrs(node, ctx); bool const keepdims = attrs.get("keepdims", 1); int32_t const ndim = input.shape().nbDims; @@ -5034,6 +5031,10 @@ DEFINE_BUILTIN_OP_IMPORTER(ReduceLogSumExp) } 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); } @@ -5043,6 +5044,9 @@ DEFINE_BUILTIN_OP_IMPORTER(ReduceLogSumExp) } std::sort(axes.begin(), axes.end()); + 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) { diff --git a/onnx_backend_test.py b/onnx_backend_test.py index 6616452b..3813afcb 100644 --- a/onnx_backend_test.py +++ b/onnx_backend_test.py @@ -201,6 +201,46 @@ def test_reduce_log_sum_exp_large_finite_input_custom(self): 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"], ["y"], keepdims=0, axes=[1]) + 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])], + ) + 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.shape, x.shape) + np.testing.assert_allclose(actual, x, rtol=1e-6, atol=1e-6) + globals().update(backend_test .enable_report() .test_cases) From d1e3a9650c3d28e4fa3a6599d781402ccd3c9156 Mon Sep 17 00:00:00 2001 From: peter941221 Date: Wed, 10 Jun 2026 12:39:38 +0800 Subject: [PATCH 3/4] Use axes input in ReduceLogSumExp test Signed-off-by: peter941221 --- onnx_backend_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/onnx_backend_test.py b/onnx_backend_test.py index 3813afcb..a8a38086 100644 --- a/onnx_backend_test.py +++ b/onnx_backend_test.py @@ -202,12 +202,13 @@ def test_reduce_log_sum_exp_large_finite_input_custom(self): self.assertAlmostEqual(actual, float(expected), places=5) def test_reduce_log_sum_exp_keepdims_zero_custom(self): - node = onnx_helper.make_node("ReduceLogSumExp", ["x"], ["y"], keepdims=0, axes=[1]) + 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 From 672f12127b0c68f00b8a4e117dbfb0078f3cb734 Mon Sep 17 00:00:00 2001 From: peter941221 Date: Wed, 10 Jun 2026 12:45:21 +0800 Subject: [PATCH 4/4] Relax noop shape assertion in backend test Signed-off-by: peter941221 --- onnx_backend_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnx_backend_test.py b/onnx_backend_test.py index a8a38086..d501e438 100644 --- a/onnx_backend_test.py +++ b/onnx_backend_test.py @@ -239,8 +239,8 @@ def test_reduce_log_sum_exp_noop_empty_axes_custom(self): outputs = trt.run_model(model, [x], device="CUDA:0") actual = np.asarray(outputs.y) - self.assertEqual(actual.shape, x.shape) - np.testing.assert_allclose(actual, x, rtol=1e-6, atol=1e-6) + 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()