diff --git a/tmva/sofie/test/TestCustomModelsFromONNX.cxx b/tmva/sofie/test/TestCustomModelsFromONNX.cxx index 94993f601a3c4..07bc0d7480d87 100644 --- a/tmva/sofie/test/TestCustomModelsFromONNX.cxx +++ b/tmva/sofie/test/TestCustomModelsFromONNX.cxx @@ -918,6 +918,21 @@ TEST(ONNX, Pow_broadcast){ } +TEST(ONNX, ReduceMean_kFirst) +{ + // ReduceMean over axis=0 (kFirst path) on a [3,4] tensor. + std::vector input(12); + std::iota(input.begin(), input.end(), 0.0f); + std::vector correct_output = {4, 5, 6, 7}; + + ASSERT_INCLUDE_AND_RUN(std::vector, "ReduceMean_kFirst", input); + + EXPECT_EQ(output.size(), correct_output.size()); + for (size_t i = 0; i < output.size(); ++i) { + EXPECT_LE(std::abs(output[i] - correct_output[i]), DEFAULT_TOLERANCE); + } +} + TEST(ONNX, ReduceProd){ constexpr float TOLERANCE = DEFAULT_TOLERANCE; diff --git a/tmva/sofie/test/input_models/ReduceMean_kFirst.onnx b/tmva/sofie/test/input_models/ReduceMean_kFirst.onnx new file mode 100644 index 0000000000000..7d5e7bad145f6 Binary files /dev/null and b/tmva/sofie/test/input_models/ReduceMean_kFirst.onnx differ