[tmva][sofie] Add support for the ONNX ArgMax operator

tmva/sofie/README.md

@@ -89,6 +89,7 @@ std::vector<std::string> supportedOperators = parser.GetRegisteredOperators();
 
 - [x] Abs
 - [x] Add
+- [x] ArgMax
 - [x] AveragePool
 - [x] BatchNormalization
 - [x] Cast

tmva/sofie/inc/TMVA/ROperator_ArgMax.hxx

@@ -0,0 +1,111 @@
+#ifndef TMVA_SOFIE_ROPERATOR_ARGMAX
+#define TMVA_SOFIE_ROPERATOR_ARGMAX
+
+#include "TMVA/SOFIE_common.hxx"
+#include "TMVA/ROperator.hxx"
+#include "TMVA/RModel.hxx"
+
+#include <sstream>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace TMVA {
+namespace Experimental {
+namespace SOFIE {
+
+class ROperator_ArgMax final : public ROperator {
+private:
+   /* Attributes */
+   int64_t fAttrAxis = 0;
+   int fKeepdims = 1;
+   int fSelectLastIndex = 0;
+
+   /* Tensor names */
+   std::string fNX; // input
+   std::string fNY; // output
+
+   /* Tensor shapes */
+   std::vector<Dim> fShapeX;
+   std::vector<Dim> fShapeY;
+
+public:
+   ROperator_ArgMax() {}
+   ROperator_ArgMax(int64_t axis, int keepdims, int selectLastIndex, std::string nameX, std::string nameY)
+      : fAttrAxis(axis),
+        fKeepdims(keepdims),
+        fSelectLastIndex(selectLastIndex),
+        fNX(UTILITY::Clean_name(nameX)),
+        fNY(UTILITY::Clean_name(nameY))
+   {
+      fInputTensorNames = {fNX};
+      fOutputTensorNames = {fNY};
+   }
+
+   void Initialize(RModel &model) override
+   {
+      if (!model.CheckIfTensorAlreadyExist(fNX)) {
+         throw std::runtime_error("TMVA SOFIE ArgMax Op Input Tensor " + fNX + " is not found in model");
+      }
+      fShapeX = model.GetDimTensorShape(fNX);
+      int64_t rank = static_cast<int64_t>(fShapeX.size());
+
+      // negative axis counts from the end (e.g. -1 is the last dimension)
+      if (fAttrAxis < 0) {
+         fAttrAxis += rank;
+      }
+      // valid range is [-rank, rank-1], i.e. [0, rank-1] after normalization
+      if (fAttrAxis < 0 || fAttrAxis >= rank) {
+         throw std::runtime_error("TMVA SOFIE ArgMax Op - invalid axis value " + std::to_string(fAttrAxis) +
+                                  " for input tensor " + fNX + " of rank " + std::to_string(rank));
+      }
+
+      // compute output shape: reduced axis becomes 1 (keepdims) or is removed
+      fShapeY = fShapeX;
+      if (fKeepdims != 0) {
+         fShapeY[fAttrAxis] = Dim{1};
+      } else {
+         fShapeY.erase(fShapeY.begin() + fAttrAxis);
+      }
+      // output contains indices, so its type is always INT64 regardless of the input type
+      model.AddIntermediateTensor(fNY, ETensorType::INT64, fShapeY);
+   }
+
+   std::string Generate(std::string opName) override
+   {
+      // split the input shape around the axis: [outer..., axisDim, inner...]
+      std::string outer = ConvertDimShapeToLength(std::vector<Dim>(fShapeX.begin(), fShapeX.begin() + fAttrAxis));
+      std::string axisDim = fShapeX[fAttrAxis].GetVal();
+      std::string inner = ConvertDimShapeToLength(std::vector<Dim>(fShapeX.begin() + fAttrAxis + 1, fShapeX.end()));
+
+      // element (o, k, i) of the input is at o * axisDim * inner + k * inner + i
+      std::string stride = "(" + axisDim + ") * (" + inner + ")";
+      // with '>' ties keep the first index, with '>=' they move to the last one
+      std::string cmp = (fSelectLastIndex != 0) ? ">=" : ">";
+
+      std::stringstream out;
+      out << "\n//------ ArgMax  " << opName << "\n";
+      out << SP << "for (size_t o = 0; o < " << outer << "; o++) {\n";
+      out << SP << SP << "for (size_t i = 0; i < " << inner << "; i++) {\n";
+      out << SP << SP << SP << "size_t base = o * " << stride << " + i;\n";
+      out << SP << SP << SP << "auto greatest = tensor_" << fNX << "[base];\n";
+      out << SP << SP << SP << "int64_t bestIndex = 0;\n";
+      out << SP << SP << SP << "for (size_t k = 1; k < " << axisDim << "; k++) {\n";
+      out << SP << SP << SP << SP << "auto value = tensor_" << fNX << "[base + k * (" << inner << ")];\n";
+      out << SP << SP << SP << SP << "if (value " << cmp << " greatest) {\n";
+      out << SP << SP << SP << SP << SP << "greatest = value;\n";
+      out << SP << SP << SP << SP << SP << "bestIndex = static_cast<int64_t>(k);\n";
+      out << SP << SP << SP << SP << "}\n";
+      out << SP << SP << SP << "}\n";
+      out << SP << SP << SP << "tensor_" << fNY << "[o * (" << inner << ") + i] = bestIndex;\n";
+      out << SP << SP << "}\n";
+      out << SP << "}\n";
+      return out.str();
+   }
+};
+
+} // namespace SOFIE
+} // namespace Experimental
+} // namespace TMVA
+
+#endif // TMVA_SOFIE_ROPERATOR_ARGMAX

tmva/sofie/test/TestCustomModelsFromONNX.cxx

@@ -615,6 +615,38 @@ TEST(ONNX, ReduceMin)
    expectNear(output, correct_output, DEFAULT_TOLERANCE);
 }
 
+TEST(ONNX, ArgMax)
+{
+   // default attributes: axis=0, keepdims=1 on a [2,3] tensor
+   std::vector<float> input({1, 9, 3, 7, 2, 8});
+   std::vector<int64_t> correct_output({1, 0, 1});
+
+   ASSERT_INCLUDE_AND_RUN(std::vector<int64_t>, "ArgMax", input);
+
+   expectEqual(output, correct_output);
+}
+
+TEST(ONNX, ArgMax_Axis1_NoKeepdims)
+{
+   std::vector<float> input({1, 9, 3, 7, 2, 8});
+   std::vector<int64_t> correct_output({1, 2});
+
+   ASSERT_INCLUDE_AND_RUN(std::vector<int64_t>, "ArgMax_Axis1_NoKeepdims", input);
+
+   expectEqual(output, correct_output);
+}
+
+TEST(ONNX, ArgMax_NegAxis_SelectLast)
+{
+   // axis=-1 with ties: select_last_index=1 picks the last maximum
+   std::vector<float> input({5, 5, 1, 2, 7, 7});
+   std::vector<int64_t> correct_output({1, 2});
+
+   ASSERT_INCLUDE_AND_RUN(std::vector<int64_t>, "ArgMax_NegAxis_SelectLast", input);
+
+   expectEqual(output, correct_output);
+}
+
 // Elu on a tensor whose first dimension is only known at run time.
 TEST(ONNX, EluDynShape)
 {

tmva/sofie/test/generate_input_models.py

@@ -245,6 +245,39 @@ def make_AddBroadcast7():
     return _model(graph, opset=17, ir_version=8)
 
 
+def _argmax_model(name, out_shape, **attrs):
+    """ArgMax on a [2, 3] float tensor with the given attributes."""
+    nodes = [
+        helper.make_node("ArgMax", ["input"], ["output"], **attrs),
+    ]
+    graph = helper.make_graph(
+        nodes,
+        name,
+        inputs=[
+            _vi("input", FLOAT, [2, 3]),
+        ],
+        outputs=[
+            _vi("output", INT64, out_shape),
+        ],
+    )
+    return _model(graph, opset=13, ir_version=8)
+
+
+def make_ArgMax():
+    """Ops: ArgMax (default attributes: axis=0, keepdims=1)"""
+    return _argmax_model("argmax_test", [1, 3])
+
+
+def make_ArgMax_Axis1_NoKeepdims():
+    """Ops: ArgMax"""
+    return _argmax_model("argmax_axis1_test", [2], axis=1, keepdims=0)
+
+
+def make_ArgMax_NegAxis_SelectLast():
+    """Ops: ArgMax"""
+    return _argmax_model("argmax_select_last_test", [2, 1], axis=-1, keepdims=1, select_last_index=1)
+
+
 def make_AveragePool1d_CeilMode_Overhang():
     """Ops: AveragePool"""
     nodes = [
@@ -6164,6 +6197,9 @@ def make_AliasDynShapeAcrossNewTensor():
     "AddBroadcast5": make_AddBroadcast5,
     "AddBroadcast6": make_AddBroadcast6,
     "AddBroadcast7": make_AddBroadcast7,
+    "ArgMax": make_ArgMax,
+    "ArgMax_Axis1_NoKeepdims": make_ArgMax_Axis1_NoKeepdims,
+    "ArgMax_NegAxis_SelectLast": make_ArgMax_NegAxis_SelectLast,
     "Asinh": make_Asinh,
     "Atanh": make_Atanh,
     "AveragePool1d_CeilMode": make_AveragePool1d_CeilMode,

tmva/sofie_parsers/CMakeLists.txt

@@ -78,6 +78,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTTMVASofieParser
     src/ParseGelu.cxx
     src/ParseHardSigmoid.cxx
     src/ParseHardSwish.cxx
+    src/ParseArgMax.cxx
   DEPENDENCIES
     ROOTTMVASofie
 )

tmva/sofie_parsers/src/ParseArgMax.cxx

@@ -0,0 +1,44 @@
+#include "TMVA/RModelParser_ONNX.hxx"
+#include "TMVA/ROperator_ArgMax.hxx"
+#include "onnx.hxx"
+
+namespace TMVA {
+namespace Experimental {
+namespace SOFIE {
+
+ParserFuncSignature ParseArgMax = [](RModelParser_ONNX &parser, const onnx::NodeProto &nodeproto) {
+   auto input_name = nodeproto.input(0);
+   if (!parser.IsRegisteredTensorType(input_name)) {
+      throw std::runtime_error("TMVA::SOFIE ONNX Parser ArgMax op has input tensor " + input_name +
+                               " but its type is not yet registered");
+   }
+
+   // default values as in the ONNX specification
+   int64_t attr_axis = 0;
+   int attr_keepdims = 1;
+   int attr_select_last_index = 0;
+   for (int i = 0; i < nodeproto.attribute_size(); i++) {
+      const std::string &attribute_name = nodeproto.attribute(i).name();
+      if (attribute_name == "axis")
+         attr_axis = nodeproto.attribute(i).i();
+      else if (attribute_name == "keepdims")
+         attr_keepdims = nodeproto.attribute(i).i();
+      else if (attribute_name == "select_last_index")
+         attr_select_last_index = nodeproto.attribute(i).i();
+   }
+
+   std::string output_name = nodeproto.output(0);
+   std::unique_ptr<ROperator> op(
+      new ROperator_ArgMax(attr_axis, attr_keepdims, attr_select_last_index, input_name, output_name));
+
+   // the output contains indices, so it is always INT64
+   if (!parser.IsRegisteredTensorType(output_name)) {
+      parser.RegisterTensorType(output_name, ETensorType::INT64);
+   }
+
+   return op;
+};
+
+} // namespace SOFIE
+} // namespace Experimental
+} // namespace TMVA

tmva/sofie_parsers/src/RModelParser_ONNX.cxx

@@ -84,6 +84,7 @@ extern ParserFuncSignature ParseScatterElements;
 extern ParserFuncSignature ParseScatterND;
 extern ParserFuncSignature ParseNonZero;
 extern ParserFuncSignature ParseClip;
+extern ParserFuncSignature ParseArgMax;
 // Declaration of fused operators
 extern ParserFuseFuncSignature ParseFuseConvAdd;
 extern ParserFuseFuncSignature ParseFuseGemmRelu;
@@ -357,6 +358,7 @@ RModelParser_ONNX::RModelParser_ONNX() noexcept : fOperatorsMapImpl(std::make_un
    RegisterOperator("ScatterND", ParseScatterND);
    RegisterOperator("NonZero", ParseNonZero);
    RegisterOperator("Clip", ParseClip);
+   RegisterOperator("ArgMax", ParseArgMax);
 }
 
 // Destructor of the parser