From eab9b5adf42dac1fa54fcbfd7dabb37fb846d2d9 Mon Sep 17 00:00:00 2001 From: ANIRUDDHA ADAK Date: Wed, 30 Sep 2026 18:33:53 +0530 Subject: [PATCH 1/2] [tmva][sofie] Fix -inf in LogSoftmax for inputs far apart from the maximum ROperator_Softmax computed LogSoftmax as log(softmax(x)): it normalized exp(x - max) by the sum and then took the log of the result. When an input is more than ~88 below the maximum, exp(x - max) underflows to 0 in float, so the log returns -inf instead of a large negative number, and values close to the underflow lose precision as well (e.g. -99.98 instead of -100). Evaluate the numerically stable form (x - max) - log(sum(exp(x - max))) instead, which does not depend on exp(x - max) staying above the underflow threshold. Both code paths of the operator are fixed: the contiguous last-axis case and the generic strided case. The plain Softmax path is unchanged. There were no LogSoftmax tests, so this adds five ONNX test models (LogSoftmax1d/2d/3d for the ordinary cases, plus LogSoftmaxLargeRange and LogSoftmaxLargeRangeAxis0, which cover both code paths with inputs 100 apart and returned -inf before). The reference outputs of the two large-range models are computed with the stable formula in double precision, because onnx.reference computes log(softmax(x)) and is wrong for the same reason. Fixes #23547 --- tmva/sofie/inc/TMVA/ROperator_Softmax.hxx | 56 +++++--- tmva/sofie/test/TestCustomModelsFromONNX.cxx | 50 ++++++++ tmva/sofie/test/generate_input_models.py | 128 +++++++++++++++++++ 3 files changed, 217 insertions(+), 17 deletions(-) diff --git a/tmva/sofie/inc/TMVA/ROperator_Softmax.hxx b/tmva/sofie/inc/TMVA/ROperator_Softmax.hxx index 025d6d678088a..e24fdfc66fd32 100644 --- a/tmva/sofie/inc/TMVA/ROperator_Softmax.hxx +++ b/tmva/sofie/inc/TMVA/ROperator_Softmax.hxx @@ -104,12 +104,21 @@ public: out << SP << SP << SP << "sum += y_ptr[j];\n"; out << SP << SP << "}\n"; - out << SP << SP << fType << " inv_sum = 1.0f / sum;\n"; - out << SP << SP << "for (int j = 0; j < " << axis_size << "; ++j) {\n"; - out << SP << SP << SP << "y_ptr[j] *= inv_sum;\n"; - if (fLogSoftmax) - out << SP << SP << SP << "y_ptr[j] = " << logFunction << "(y_ptr[j]);\n"; - out << SP << SP << "}\n"; + if (fLogSoftmax) { + // evaluate (x - vmax) - log(sum(exp(x - vmax))) instead of + // log(exp(x - vmax) / sum): for inputs far apart from the maximum + // exp(x - vmax) underflows to 0 and its log is -inf, while the + // correct result is a large negative number (see #23547) + out << SP << SP << fType << " log_sum = " << logFunction << "(sum);\n"; + out << SP << SP << "for (int j = 0; j < " << axis_size << "; ++j) {\n"; + out << SP << SP << SP << "y_ptr[j] = (x_ptr[j] - vmax) - log_sum;\n"; + out << SP << SP << "}\n"; + } else { + out << SP << SP << fType << " inv_sum = 1.0f / sum;\n"; + out << SP << SP << "for (int j = 0; j < " << axis_size << "; ++j) {\n"; + out << SP << SP << SP << "y_ptr[j] *= inv_sum;\n"; + out << SP << SP << "}\n"; + } out << SP << "}\n"; } else { @@ -166,20 +175,33 @@ public: for (size_t j = 0; j < size-1; j++) out << SP; out << "}\n"; // normalize - for (size_t j = 0; j < size-1; j++) out << SP; - out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n"; - for (size_t j = 0; j < size; j++) out << SP; - out << "size_t id = index + i"; - if (stride[axis].GetVal() != "1") out << "*(" << stride[axis] << ")"; - out << ";\n"; - for (size_t j = 0; j < size; j++) out << SP; - out << "tensor_" << fNY << "[id] /= sum;\n"; if (fLogSoftmax) { + // same numerically stable form as in the last-axis case above: + // (x - vmax) - log(sum(exp(x - vmax))), see #23547 + for (size_t j = 0; j < size-1; j++) out << SP; + out << fType << " log_sum = " << logFunction << "(sum);\n"; + for (size_t j = 0; j < size-1; j++) out << SP; + out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n"; + for (size_t j = 0; j < size; j++) out << SP; + out << "size_t id = index + i"; + if (stride[axis].GetVal() != "1") out << "*(" << stride[axis] << ")"; + out << ";\n"; + for (size_t j = 0; j < size; j++) out << SP; + out << "tensor_" << fNY << "[id] = (tensor_" << fNX << "[id] - vmax) - log_sum;\n"; + for (size_t j = 0; j < size-1; j++) out << SP; + out << "}\n"; + } else { + for (size_t j = 0; j < size-1; j++) out << SP; + out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n"; + for (size_t j = 0; j < size; j++) out << SP; + out << "size_t id = index + i"; + if (stride[axis].GetVal() != "1") out << "*(" << stride[axis] << ")"; + out << ";\n"; for (size_t j = 0; j < size; j++) out << SP; - out << "tensor_" << fNY << "[id] = " << logFunction << "(tensor_" << fNY << "[id]);\n"; + out << "tensor_" << fNY << "[id] /= sum;\n"; + for (size_t j = 0; j < size-1; j++) out << SP; + out << "}\n"; } - for (size_t j = 0; j < size-1; j++) out << SP; - out << "}\n"; //end loops for (int i = static_cast(k) - 1; i >= 0; i--) { for (int j = 0; j < i; j++) out << SP; diff --git a/tmva/sofie/test/TestCustomModelsFromONNX.cxx b/tmva/sofie/test/TestCustomModelsFromONNX.cxx index ad1f7286d4555..9ad1ff693a838 100644 --- a/tmva/sofie/test/TestCustomModelsFromONNX.cxx +++ b/tmva/sofie/test/TestCustomModelsFromONNX.cxx @@ -1024,6 +1024,56 @@ TEST(ONNX, Softmax4d) expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); } +TEST(ONNX, LogSoftmax1d) +{ + SofieReference ref = readReference("LogSoftmax1d"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "LogSoftmax1d", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + +TEST(ONNX, LogSoftmax2d) +{ + SofieReference ref = readReference("LogSoftmax2d"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "LogSoftmax2d", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + +TEST(ONNX, LogSoftmax3d) +{ + SofieReference ref = readReference("LogSoftmax3d"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "LogSoftmax3d", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + +// Inputs more than ~88 apart from the maximum make exp(x - max) underflow to 0 +// in float32, so log(softmax(x)) returns -inf where the stable form +// (x - max) - log(sum(exp(x - max))) returns a large negative number +TEST(ONNX, LogSoftmaxLargeRange) +{ + SofieReference ref = readReference("LogSoftmaxLargeRange"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "LogSoftmaxLargeRange", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + +// same as above, but reducing along the first axis, which takes the generic +// code path of the SOFIE Softmax operator +TEST(ONNX, LogSoftmaxLargeRangeAxis0) +{ + SofieReference ref = readReference("LogSoftmaxLargeRangeAxis0"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "LogSoftmaxLargeRangeAxis0", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + TEST(ONNX, ConvTranspose1d) { SofieReference ref = readReference("ConvTranspose1d"); diff --git a/tmva/sofie/test/generate_input_models.py b/tmva/sofie/test/generate_input_models.py index 28a06d4ae68eb..1668eeff5dcef 100644 --- a/tmva/sofie/test/generate_input_models.py +++ b/tmva/sofie/test/generate_input_models.py @@ -3532,6 +3532,103 @@ def make_Log(): return _model(graph, opset=14, ir_version=7, producer_name='pytorch', producer_version='1.13.1') +def make_LogSoftmax1d(): + """Ops: LogSoftmax""" + nodes = [ + helper.make_node('LogSoftmax', ['X'], ['Y']), + ] + graph = helper.make_graph( + nodes, + 'LogSoftmax', + inputs=[ + _vi('X', FLOAT, [3]), + ], + outputs=[ + _vi('Y', FLOAT, [3]), + ], + ) + return _model(graph, opset=17, ir_version=8) + + +def make_LogSoftmax2d(): + """Ops: LogSoftmax""" + nodes = [ + helper.make_node('LogSoftmax', ['X'], ['Y'], axis=-1), + ] + graph = helper.make_graph( + nodes, + 'LogSoftmax', + inputs=[ + _vi('X', FLOAT, [2, 3]), + ], + outputs=[ + _vi('Y', FLOAT, [2, 3]), + ], + ) + return _model(graph, opset=17, ir_version=8) + + +def make_LogSoftmax3d(): + """Ops: LogSoftmax""" + nodes = [ + helper.make_node('LogSoftmax', ['X'], ['Y'], axis=1), + ] + graph = helper.make_graph( + nodes, + 'LogSoftmax', + inputs=[ + _vi('X', FLOAT, [2, 3, 4]), + ], + outputs=[ + _vi('Y', FLOAT, [2, 3, 4]), + ], + ) + return _model(graph, opset=17, ir_version=8) + + +def make_LogSoftmaxLargeRange(): + """Ops: LogSoftmax + + The inputs of each row are 100 apart, so exp(x - max) underflows to 0 in + float32 for all but the largest one. log(softmax(x)) is -inf there, while + the stable form (x - max) - log(sum(exp(x - max))) stays finite.""" + nodes = [ + helper.make_node('LogSoftmax', ['X'], ['Y'], axis=-1), + ] + graph = helper.make_graph( + nodes, + 'LogSoftmax', + inputs=[ + _vi('X', FLOAT, [2, 3]), + ], + outputs=[ + _vi('Y', FLOAT, [2, 3]), + ], + ) + return _model(graph, opset=17, ir_version=8) + + +def make_LogSoftmaxLargeRangeAxis0(): + """Ops: LogSoftmax + + Same as LogSoftmaxLargeRange but with the reduction along the first axis, + which takes the generic code path of the SOFIE Softmax operator.""" + nodes = [ + helper.make_node('LogSoftmax', ['X'], ['Y'], axis=0), + ] + graph = helper.make_graph( + nodes, + 'LogSoftmax', + inputs=[ + _vi('X', FLOAT, [3, 2]), + ], + outputs=[ + _vi('Y', FLOAT, [3, 2]), + ], + ) + return _model(graph, opset=17, ir_version=8) + + def make_MatMul_1D_Constant(): """Ops: MatMul""" nodes = [ @@ -5955,6 +6052,11 @@ def make_IdentityWeightOutput(): "Linear_32": make_Linear_32, "Linear_64": make_Linear_64, "Log": make_Log, + "LogSoftmax1d": make_LogSoftmax1d, + "LogSoftmax2d": make_LogSoftmax2d, + "LogSoftmax3d": make_LogSoftmax3d, + "LogSoftmaxLargeRange": make_LogSoftmaxLargeRange, + "LogSoftmaxLargeRangeAxis0": make_LogSoftmaxLargeRangeAxis0, "MatMul_1D_Constant": make_MatMul_1D_Constant, "MatMul_Stacked": make_MatMul_Stacked, "MatMul_Stacked2": make_MatMul_Stacked2, @@ -6862,6 +6964,13 @@ def rand_f32(seed, shape): "Linear_32": [f32(np.full(3200, 1.0), (32, 100))], "Linear_64": [f32(np.full(6400, 1.0), (64, 100))], "Log": [f32([1.0, 2.0, 3.0, 4.0], (4,))], + "LogSoftmax1d": [f32([-1.0, 0.0, 1.0], (3,))], + "LogSoftmax2d": [f32([-1.0, 0.0, 1.0, 2.0, 1.0, 0.0], (2, 3))], + "LogSoftmax3d": [rand_f32(11, (2, 3, 4))], + # 100 apart per row: exp(x - max) underflows to 0 in float32 away from the + # maximum, which log(softmax(x)) turns into -inf + "LogSoftmaxLargeRange": [f32([0.0, 100.0, 200.0, 200.0, 100.0, 0.0], (2, 3))], + "LogSoftmaxLargeRangeAxis0": [f32([0.0, 100.0, 100.0, 200.0, 200.0, 300.0], (3, 2))], "Max": [ f32([1.0, 2.0, -1.0], (1, 3)), f32([3.0, 0.0, 4.0], (1, 3)), @@ -7698,6 +7807,23 @@ def resolve(name): return [y.astype(np.float32)] +def _logsoftmax_reference(model, feeds): + """LogSoftmax evaluated with the numerically stable formula in double + precision. + + The ReferenceEvaluator computes log(softmax(x)), whose exp(x - max) + underflows to 0 in float32 once the inputs are more than ~88 apart from the + maximum, so it yields -inf where the correct result is a large negative + number (see root-project/root#23547).""" + node = next(n for n in model.graph.node if n.op_type == "LogSoftmax") + axis = next((a.i for a in node.attribute if a.name == "axis"), -1) + x = np.asarray(feeds[node.input[0]], dtype=np.float64) + axis = axis % x.ndim + shifted = x - x.max(axis=axis, keepdims=True) + y = shifted - np.log(np.exp(shifted).sum(axis=axis, keepdims=True)) + return [y.astype(np.float32)] + + EXPECTED_OVERRIDES = { "GRUBidirectional": _recurrent_reference, "LSTMBidirectional": _recurrent_reference, @@ -7712,6 +7838,8 @@ def resolve(name): "IdentityWeightBatchNorm": _batchnorm_reference, "BatchNormEpsilon": _batchnorm_reference, "BatchNormReluEpsilon": _batchnorm_reference, + "LogSoftmaxLargeRange": _logsoftmax_reference, + "LogSoftmaxLargeRangeAxis0": _logsoftmax_reference, } From c99473b5c08b417004990977a8c143cbe785ba27 Mon Sep 17 00:00:00 2001 From: ANIRUDDHA ADAK Date: Wed, 30 Sep 2026 18:55:47 +0530 Subject: [PATCH 2/2] [tmva][sofie] Apply ROOT's clang-format and ruff formatting to the LogSoftmax fix The formatting checks flagged the lines added by the previous commit: - clang-format 20 wants the one-line loop bodies split and `size - 1` spaced, so the new generic-path block is now written the way .clang-format formats those lines. Only the host code is reformatted; the inference code that ROperator_Softmax emits is unchanged. - ruff format, which the CI runs with --preview over the changed ranges, wants double quotes, so the new make_LogSoftmax* builders use them. The surrounding functions in generate_input_models.py still use single quotes, but they lie outside the changed ranges the check covers and are left untouched. Verified with the same commands the CI runs: clang-format-20 restricted to the diff hunks via --lines, and ruff format --diff --preview --range over the Python hunks. Both are clean now. --- tmva/sofie/inc/TMVA/ROperator_Softmax.hxx | 33 ++++++++++++------- tmva/sofie/test/generate_input_models.py | 40 +++++++++++------------ 2 files changed, 42 insertions(+), 31 deletions(-) diff --git a/tmva/sofie/inc/TMVA/ROperator_Softmax.hxx b/tmva/sofie/inc/TMVA/ROperator_Softmax.hxx index e24fdfc66fd32..7597ac1846cfe 100644 --- a/tmva/sofie/inc/TMVA/ROperator_Softmax.hxx +++ b/tmva/sofie/inc/TMVA/ROperator_Softmax.hxx @@ -178,28 +178,39 @@ public: if (fLogSoftmax) { // same numerically stable form as in the last-axis case above: // (x - vmax) - log(sum(exp(x - vmax))), see #23547 - for (size_t j = 0; j < size-1; j++) out << SP; + for (size_t j = 0; j < size - 1; j++) + out << SP; out << fType << " log_sum = " << logFunction << "(sum);\n"; - for (size_t j = 0; j < size-1; j++) out << SP; + for (size_t j = 0; j < size - 1; j++) + out << SP; out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n"; - for (size_t j = 0; j < size; j++) out << SP; + for (size_t j = 0; j < size; j++) + out << SP; out << "size_t id = index + i"; - if (stride[axis].GetVal() != "1") out << "*(" << stride[axis] << ")"; + if (stride[axis].GetVal() != "1") + out << "*(" << stride[axis] << ")"; out << ";\n"; - for (size_t j = 0; j < size; j++) out << SP; + for (size_t j = 0; j < size; j++) + out << SP; out << "tensor_" << fNY << "[id] = (tensor_" << fNX << "[id] - vmax) - log_sum;\n"; - for (size_t j = 0; j < size-1; j++) out << SP; + for (size_t j = 0; j < size - 1; j++) + out << SP; out << "}\n"; } else { - for (size_t j = 0; j < size-1; j++) out << SP; + for (size_t j = 0; j < size - 1; j++) + out << SP; out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n"; - for (size_t j = 0; j < size; j++) out << SP; + for (size_t j = 0; j < size; j++) + out << SP; out << "size_t id = index + i"; - if (stride[axis].GetVal() != "1") out << "*(" << stride[axis] << ")"; + if (stride[axis].GetVal() != "1") + out << "*(" << stride[axis] << ")"; out << ";\n"; - for (size_t j = 0; j < size; j++) out << SP; + for (size_t j = 0; j < size; j++) + out << SP; out << "tensor_" << fNY << "[id] /= sum;\n"; - for (size_t j = 0; j < size-1; j++) out << SP; + for (size_t j = 0; j < size - 1; j++) + out << SP; out << "}\n"; } //end loops diff --git a/tmva/sofie/test/generate_input_models.py b/tmva/sofie/test/generate_input_models.py index 1668eeff5dcef..7a84fe589997d 100644 --- a/tmva/sofie/test/generate_input_models.py +++ b/tmva/sofie/test/generate_input_models.py @@ -3535,16 +3535,16 @@ def make_Log(): def make_LogSoftmax1d(): """Ops: LogSoftmax""" nodes = [ - helper.make_node('LogSoftmax', ['X'], ['Y']), + helper.make_node("LogSoftmax", ["X"], ["Y"]), ] graph = helper.make_graph( nodes, - 'LogSoftmax', + "LogSoftmax", inputs=[ - _vi('X', FLOAT, [3]), + _vi("X", FLOAT, [3]), ], outputs=[ - _vi('Y', FLOAT, [3]), + _vi("Y", FLOAT, [3]), ], ) return _model(graph, opset=17, ir_version=8) @@ -3553,16 +3553,16 @@ def make_LogSoftmax1d(): def make_LogSoftmax2d(): """Ops: LogSoftmax""" nodes = [ - helper.make_node('LogSoftmax', ['X'], ['Y'], axis=-1), + helper.make_node("LogSoftmax", ["X"], ["Y"], axis=-1), ] graph = helper.make_graph( nodes, - 'LogSoftmax', + "LogSoftmax", inputs=[ - _vi('X', FLOAT, [2, 3]), + _vi("X", FLOAT, [2, 3]), ], outputs=[ - _vi('Y', FLOAT, [2, 3]), + _vi("Y", FLOAT, [2, 3]), ], ) return _model(graph, opset=17, ir_version=8) @@ -3571,16 +3571,16 @@ def make_LogSoftmax2d(): def make_LogSoftmax3d(): """Ops: LogSoftmax""" nodes = [ - helper.make_node('LogSoftmax', ['X'], ['Y'], axis=1), + helper.make_node("LogSoftmax", ["X"], ["Y"], axis=1), ] graph = helper.make_graph( nodes, - 'LogSoftmax', + "LogSoftmax", inputs=[ - _vi('X', FLOAT, [2, 3, 4]), + _vi("X", FLOAT, [2, 3, 4]), ], outputs=[ - _vi('Y', FLOAT, [2, 3, 4]), + _vi("Y", FLOAT, [2, 3, 4]), ], ) return _model(graph, opset=17, ir_version=8) @@ -3593,16 +3593,16 @@ def make_LogSoftmaxLargeRange(): float32 for all but the largest one. log(softmax(x)) is -inf there, while the stable form (x - max) - log(sum(exp(x - max))) stays finite.""" nodes = [ - helper.make_node('LogSoftmax', ['X'], ['Y'], axis=-1), + helper.make_node("LogSoftmax", ["X"], ["Y"], axis=-1), ] graph = helper.make_graph( nodes, - 'LogSoftmax', + "LogSoftmax", inputs=[ - _vi('X', FLOAT, [2, 3]), + _vi("X", FLOAT, [2, 3]), ], outputs=[ - _vi('Y', FLOAT, [2, 3]), + _vi("Y", FLOAT, [2, 3]), ], ) return _model(graph, opset=17, ir_version=8) @@ -3614,16 +3614,16 @@ def make_LogSoftmaxLargeRangeAxis0(): Same as LogSoftmaxLargeRange but with the reduction along the first axis, which takes the generic code path of the SOFIE Softmax operator.""" nodes = [ - helper.make_node('LogSoftmax', ['X'], ['Y'], axis=0), + helper.make_node("LogSoftmax", ["X"], ["Y"], axis=0), ] graph = helper.make_graph( nodes, - 'LogSoftmax', + "LogSoftmax", inputs=[ - _vi('X', FLOAT, [3, 2]), + _vi("X", FLOAT, [3, 2]), ], outputs=[ - _vi('Y', FLOAT, [3, 2]), + _vi("Y", FLOAT, [3, 2]), ], ) return _model(graph, opset=17, ir_version=8)