Merge branch 'fix/julien/more-onnx-ops' into 'master'

Add more ONNX operators to NIER

Closes #21 and #22

See merge request laiser/caisar!50

lib/ir/nier_cfg.ml

@@ -125,30 +125,42 @@ module Node = struct
 
   type operator =
     | Add
+    | Sub
     | Mul
+    | Div
     | Matmul
+    | Gemm
     | LogSoftmax
     | ReLu
     | Transpose
     | Squeeze
     | MaxPool
     | Conv
+    | Reshape
+    | Flatten
     | Identity
+    | Constant
     | NO_OP
     | RW_Linearized_ReLu
 
   let str_op o =
     match o with
     | Add -> "Add"
+    | Sub -> "Sub"
     | Mul -> "Mul"
+    | Div -> "Div"
     | Matmul -> "Matmul"
+    | Gemm -> "Gemm"
     | LogSoftmax -> "LogSoftmax"
     | ReLu -> "ReLu"
     | Transpose -> "Transpose"
     | Squeeze -> "Squeeze"
     | MaxPool -> "MaxPool"
     | Conv -> "Conv"
+    | Reshape -> "Reshape"
+    | Flatten -> "Flatten"
     | Identity -> "Identity"
+    | Constant -> "Constant"
     | NO_OP -> "NO_OP"
     | RW_Linearized_ReLu -> "RW_Linearized_ReLu"
 

lib/ir/nier_cfg.mli

@@ -80,15 +80,21 @@ module Node : sig
 
   type operator =
     | Add
+    | Sub
     | Mul
+    | Div
     | Matmul
+    | Gemm
     | LogSoftmax
     | ReLu
     | Transpose
     | Squeeze
     | MaxPool
     | Conv
+    | Reshape
+    | Flatten
     | Identity
+    | Constant
     | NO_OP
     | RW_Linearized_ReLu
 

lib/onnx/onnx.ml

@@ -110,14 +110,20 @@ let produce_cfg (g : Oproto.Onnx.GraphProto.t) =
       | Some o -> (
         match o with
         | "Add" -> NCFG.Node.Add
+        | "Sub" -> NCFG.Node.Sub
         | "Mul" -> NCFG.Node.Mul
+        | "Div" -> NCFG.Node.Div
         | "Relu" -> NCFG.Node.ReLu
         | "MatMul" -> NCFG.Node.Matmul
+        | "Gemm" -> NCFG.Node.Gemm
         | "LogSoftmax" -> NCFG.Node.LogSoftmax
         | "Transpose" -> NCFG.Node.Transpose
         | "Squeeze" -> NCFG.Node.Squeeze
         | "MaxPool" -> NCFG.Node.MaxPool
+        | "Constant" -> NCFG.Node.Constant
         | "Conv" -> NCFG.Node.Conv
+        | "Reshape" -> NCFG.Node.Reshape
+        | "Flatten" -> NCFG.Node.Flatten
         | "Identity" -> NCFG.Node.Identity
         | _ -> raise (ParseError ("Unsupported ONNX operator " ^ o)))
     in