Skip to content

Use a stable lowering for ReduceLogSumExp - #1055

Open
peter941221 wants to merge 4 commits into
onnx:mainfrom
peter941221:fix/reduce-log-sum-exp-stability
Open

Use a stable lowering for ReduceLogSumExp#1055
peter941221 wants to merge 4 commits into
onnx:mainfrom
peter941221:fix/reduce-log-sum-exp-stability

Conversation

@peter941221

@peter941221 peter941221 commented Jun 10, 2026

Copy link
Copy Markdown

Refs NVIDIA/TensorRT#4772

Problem
The ReduceLogSumExp importer exponentiates the input before reduction. Large finite inputs overflow in that first exp, so TensorRT returns inf where ONNX Runtime stays finite.

Change
Use the stable form max(x) + log(sum(exp(x - max(x)))) for ReduceLogSumExp lowering.
Reduce the max over the same axes as the sum.
Unsqueeze the reduced max when keepdims=0 so the subtraction still broadcasts correctly.
Preserve noop_with_empty_axes=1 when opset 18 passes an empty axes input.
Extend the regression coverage for large finite inputs, keepdims=0 with axes input, and noop_with_empty_axes=1 with an empty axes input.

Validation
I reproduced NVIDIA/TensorRT#4772 with TensorRT 10.16.1.11.
Input: [250.0, 248.0, 255.0, 251.0]
ONNX Runtime: 255.025634765625
TensorRT before: inf
TensorRT after patched parser: 255.025634765625

I rebuilt the parser against the 10.16 source tree and validated it by temporarily swapping the parser library in a TensorRT 10.16.1.11 Python environment.
python -m unittest -v onnx_backend_test.TensorRTCustomReduceLogSumExpTest

Signed-off-by: peter941221 <peter941221@gmail.com>
Signed-off-by: peter941221 <peter941221@gmail.com>
Signed-off-by: peter941221 <peter941221@gmail.com>
Signed-off-by: peter941221 <peter941221@gmail.com>
@peter941221
peter941221 marked this pull request as ready for review June 10, 2026 04:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant