[tmva][sofie] Improvements to SOFIE Keras parser on the Python side

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/_gnn.py

@@ -133,21 +133,20 @@ def add_weights(gin, weights, function_target):
         function_target: Target for the function to update either of nodes, edges or globals
 
     """
-    import ROOT
+    from ROOT.TMVA.Experimental import SOFIE
 
-    if function_target == ROOT.TMVA.Experimental.SOFIE.FunctionTarget.NODES:
+    if function_target == SOFIE.FunctionTarget.NODES:
         model_block = gin.nodes_update_block
-    elif function_target == ROOT.TMVA.Experimental.SOFIE.FunctionTarget.EDGES:
+    elif function_target == SOFIE.FunctionTarget.EDGES:
         model_block = gin.edges_update_block
     else:
         model_block = gin.globals_update_block
 
     for i in weights:
-        shape = ROOT.std.vector['std::size_t']()
-        shape_as_list = i.shape.as_list()
-        for j in shape_as_list:
-            shape.push_back(j)
-        model_block.GetFunctionBlock().AddInitializedTensor['float'](i.name, shape, i.numpy())
+        model_block.GetFunctionBlock().AddInitializedTensor(
+            i.name, SOFIE.ETensorType.FLOAT, i.shape.as_list(), i.numpy()
+        )
+
 
 def add_aggregate_function(gin, reducer, relation):
     """

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/_sofie/_parser/_keras/layers/binary.py

@@ -1,23 +1,5 @@
 def MakeKerasBinary(layer):
     from ROOT.TMVA.Experimental import SOFIE
 
-    input = layer["layerInput"]
-    output = layer["layerOutput"]
-    fLayerType = layer["layerType"]
-    fLayerDType = layer["layerDType"]
-    fX1 = input[0]
-    fX2 = input[1]
-    fY = output[0]
-    op = None
-    if SOFIE.ConvertStringToType(fLayerDType) == SOFIE.ETensorType.FLOAT:
-        if fLayerType == "Add":
-            op = SOFIE.ROperator_BasicBinary(float, SOFIE.EBasicBinaryOperator.Add)(fX1, fX2, fY)
-        elif fLayerType == "Subtract":
-            op = SOFIE.ROperator_BasicBinary(float, SOFIE.EBasicBinaryOperator.Sub)(fX1, fX2, fY)
-        else:
-            op = SOFIE.ROperator_BasicBinary(float, SOFIE.EBasicBinaryOperator.Mul)(fX1, fX2, fY)
-    else:
-        raise RuntimeError(
-            "TMVA::SOFIE - Unsupported - Operator BasicBinary does not yet support input type " + fLayerDType
-        )
-    return op
+    inpt = layer["layerInput"]
+    return SOFIE.createBasicBinary(layer["layerDType"], layer["layerType"], inpt[0], inpt[1], layer["layerOutput"][0])

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/_sofie/_parser/_keras/parser.py

@@ -99,6 +99,11 @@ def move_operator(op):
         """
         import ROOT
 
+        # If the object is already held by a smart pointer, just move it.
+        smartptr = op.__smartptr__()
+        if smartptr:
+            return type(smartptr)(ROOT.std.move(smartptr))
+
         ROOT.SetOwnership(op, False)
         return ROOT.std.unique_ptr[type(op)](op)
 
@@ -131,10 +136,7 @@ def move_operator(op):
         # in c++ that does the conversion from a regular pointer to unique one in c++
         # print('adding initialized tensor..',LayerName, TargetShape)
         shape_tensor_name = LayerName + "_shape"
-        shape_data = TargetShape.data
-        print(TargetShape, shape_data)
-        print(len(TargetShape))
-        rmodel.AddInitializedTensor["int64_t"](shape_tensor_name, [len(TargetShape)], shape_data)
+        rmodel.AddInitializedTensor(shape_tensor_name, SOFIE.ETensorType.UINT64, [len(TargetShape)], TargetShape)
 
     # These layers only have one operator - excluding the recurrent layers, in which the activation function(s)
     # are included in the recurrent operator
@@ -507,7 +509,7 @@ def Parse(filename, batch_size=1):  # If a model does not have a defined batch s
 
                 else:
                     fData = fWeightArray.flatten()
-                rmodel.AddInitializedTensor["float"](fWeightName, fWeightTensorShape, fData)
+                rmodel.AddInitializedTensor(fWeightName, SOFIE.ETensorType.FLOAT, fWeightTensorShape, fData)
             else:
                 raise TypeError("Type error: TMVA SOFIE does not yet support data layer type: " + fWeightDType)
 

tmva/sofie/inc/TMVA/RModel.hxx

@@ -76,6 +76,8 @@ public:
    void AddOperator(std::unique_ptr<ROperator> op, int order_execution = -1);
    void AddInitializedTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape,
                              std::shared_ptr<void> data);
+   void AddInitializedTensor(const std::string &tensor_name, ETensorType tensor_type,
+                             const std::vector<std::size_t> &shape, void *raw_data);
    void AddConstantTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape,
                              std::shared_ptr<void> data);
 
@@ -99,15 +101,6 @@ public:
       AddConstantTensor(name, GetTemplatedType<T>(T()), shape, data_ptr);
    }
 
-   template <typename T>
-   void AddInitializedTensor(const std::string & tensor_name, const std::vector<std::size_t> & shape, T *raw_data)
-   {
-      size_t size = ConvertShapeToLength(shape);
-      std::shared_ptr<void> data(malloc(size * sizeof(T)), free);
-      std::memcpy(data.get(), raw_data, size * sizeof(T));
-      AddInitializedTensor(tensor_name,  GetTemplatedType(T()), shape, data);
-   }
-
    void AddShapeTensor(const std::string & name, const std::vector<Dim> & shapeValues, bool scalar = false);
 
    void AddExtraCodeForDimShapes(const std::string & code) { fExtraCodeForDimShapes += code; }

tmva/sofie/inc/TMVA/ROperator_BasicBinary.hxx

@@ -455,6 +455,24 @@ public:
    }
 };
 
+inline std::unique_ptr<ROperator> createBasicBinary(std::string layerDType, std::string layerType, std::string nameA,
+                                                    std::string nameB, std::string nameY)
+{
+   if (ConvertStringToType(layerDType) != ETensorType::FLOAT) {
+      throw std::runtime_error(
+         ("TMVA::SOFIE - Unsupported - Operator BasicBinary does not yet support input type " + layerDType).c_str());
+   }
+   if (layerType == "Add")
+      return std::make_unique<ROperator_BasicBinary<float, EBasicBinaryOperator::Add>>(nameA, nameB, nameY);
+   if (layerType == "Subtract")
+      return std::make_unique<ROperator_BasicBinary<float, EBasicBinaryOperator::Sub>>(nameA, nameB, nameY);
+   if (layerType == "Multiply")
+      return std::make_unique<ROperator_BasicBinary<float, EBasicBinaryOperator::Mul>>(nameA, nameB, nameY);
+
+   throw std::runtime_error(
+      ("TMVA::SOFIE - Unsupported - Operator BasicBinary does not yet support layer type " + layerType).c_str());
+}
+
 } // namespace SOFIE
 } // namespace Experimental
 } // namespace TMVA

tmva/sofie/src/RModel.cxx

@@ -229,6 +229,16 @@ void RModel::AddInitializedTensor(std::string tensor_name, ETensorType type, std
     fInitializedTensors[tensor_name] = new_tensor;
 }
 
+void RModel::AddInitializedTensor(const std::string &tensor_name, ETensorType tensor_type,
+                                  const std::vector<std::size_t> &shape, void *raw_data)
+{
+   size_t size = ConvertShapeToLength(shape);
+   auto itemsize = GetTypeSize(tensor_type);
+   std::shared_ptr<void> data(malloc(size * itemsize), free);
+   std::memcpy(data.get(), raw_data, size * itemsize);
+   AddInitializedTensor(tensor_name, tensor_type, shape, data);
+}
+
 void RModel::AddConstantTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape, std::shared_ptr<void> data) {
     tensor_name = UTILITY::Clean_name(tensor_name);
     //NB: own data