Fix argmin/max axis attribute, clean up, fix constant op in fuseOps

Author: EmergentOrder

backends/.jvm/src/main/scala/ORTOperatorBackend.scala

@@ -123,7 +123,7 @@ trait ORTOperatorBackend
     if(cacheValues.size == 0) return ModelProto()
     val nodes = cacheValues.map(_.getGraph.node).fold(Seq[NodeProto]())( (x,y) => x ++ y )
     val nodeOutputs = cacheValues.map(_.getGraph.output).fold(Seq[ValueInfoProto]())( (x,y) => x ++ y ).map(_.getName)
-    val inputs = cacheValues.map(_.getGraph.input).fold(Seq[ValueInfoProto]())( (x,y) => x ++ y ).filter(z => ! nodeOutputs.contains(z.getName)).distinct
+    val inputs = cacheValues.map(_.getGraph).filter(! _.node(0).opType.equals(Some("Constant"))).map(_.input).fold(Seq[ValueInfoProto]())( (x,y) => x ++ y ).filter(z => ! nodeOutputs.contains(z.getName)).distinct
     val outputs = cacheValues(cacheValues.size - 1).getGraph.output
     val modelProto = (cacheValues.head.clearGraph).withGraph((new GraphProto).withNode(nodes).withInput(inputs).withOutput(outputs))
     sessionCache.clear

core/src/main/scala/ONNX.scala

@@ -123,7 +123,7 @@ package object onnx {
         keepdims: KeepDims = true,
         data: Tensor[T, Tuple3[Tt, Td, S]]
     )(using tt: ValueOf[Tt1], td: TensorShapeDenotationOf[KeepOrReduceDimDenotations[Td,Axis,KeepDims]], s: ShapeOf[KeepOrReduceDims[S,Axis,KeepDims]], i: IndicesOf[Axis], k: ValueOf[KeepDims]): Tensor[Long, Tuple3[Tt1, KeepOrReduceDimDenotations[Td,Axis,KeepDims], KeepOrReduceDims[S,Axis,KeepDims]]] = {
-      val map: Map[String, Any] = Map("axis" -> indicesOf[Axis].indices.toArray, "keepdims" -> (if(valueOf[KeepDims]) 1 else 0))
+      val map: Map[String, Any] = Map("axis" -> indicesOf[Axis].indices.toArray.head, "keepdims" -> (if(valueOf[KeepDims]) 1 else 0))
       val allInputs             = Tuple1(data)
       (callOp(name, "ArgMax", allInputs, map))
     }
@@ -1217,7 +1217,7 @@ package object onnxruntime {
         selectLastIndex: Int = 0,
         data: Tensor[T, Tuple3[Tt, Td, S]]
     )(using tt: ValueOf[Tt1], td: TensorShapeDenotationOf[KeepOrReduceDimDenotations[Td,Axis,KeepDims]], s: ShapeOf[KeepOrReduceDims[S,Axis,KeepDims]], i: IndicesOf[Axis], k: ValueOf[KeepDims]): Tensor[Long, Tuple3[Tt1, KeepOrReduceDimDenotations[Td,Axis,KeepDims], KeepOrReduceDims[S,Axis,KeepDims]]] = {
-      val map: Map[String, Any] = Map("axis" -> indicesOf[Axis].indices.toArray, "select_last_index" -> selectLastIndex, "keepdims" -> (if(valueOf[KeepDims]) 1 else 0))
+      val map: Map[String, Any] = Map("axis" -> indicesOf[Axis].indices.toArray.head, "select_last_index" -> selectLastIndex, "keepdims" -> (if(valueOf[KeepDims]) 1 else 0))
       val allInputs             = Tuple1(data)
       (callOp(name, "ArgMax", allInputs, map))
     }
@@ -1234,7 +1234,7 @@ package object onnxruntime {
         selectLastIndex: Int = 0,
         data: Tensor[T, Tuple3[Tt, Td, S]]
     )(using tt: ValueOf[Tt1], td: TensorShapeDenotationOf[KeepOrReduceDimDenotations[Td,Axis,KeepDims]], s: ShapeOf[KeepOrReduceDims[S,Axis,KeepDims]], i: IndicesOf[Axis], k: ValueOf[KeepDims]): Tensor[Long, Tuple3[Tt1, KeepOrReduceDimDenotations[Td,Axis,KeepDims], KeepOrReduceDims[S,Axis,KeepDims]]] = {
-      val map: Map[String, Any] = Map("axis" -> indicesOf[Axis].indices.toArray, "select_last_index" -> selectLastIndex, "keepdims" -> (if(valueOf[KeepDims]) 1 else 0))
+      val map: Map[String, Any] = Map("axis" -> indicesOf[Axis].indices.toArray.head, "select_last_index" -> selectLastIndex, "keepdims" -> (if(valueOf[KeepDims]) 1 else 0))
       val allInputs             = Tuple1(data)
       (callOp(name, "ArgMin", allInputs, map))
     }
@@ -1249,7 +1249,7 @@ package object onnxruntime {
         keepdims: KeepDims = true,
         data: Tensor[T, Tuple3[Tt, Td, S]]
     )(using tt: ValueOf[Tt1], td: TensorShapeDenotationOf[KeepOrReduceDimDenotations[Td,Axis,KeepDims]], s: ShapeOf[KeepOrReduceDims[S,Axis,KeepDims]], i: IndicesOf[Axis], k: ValueOf[KeepDims]): Tensor[Long, Tuple3[Tt1, KeepOrReduceDimDenotations[Td,Axis,KeepDims], KeepOrReduceDims[S,Axis,KeepDims]]] = {
-      val map: Map[String, Any] = Map("axis" -> indicesOf[Axis].indices.toArray, "keepdims" -> (if(valueOf[KeepDims]) 1 else 0))
+      val map: Map[String, Any] = Map("axis" -> indicesOf[Axis].indices.toArray.head, "keepdims" -> (if(valueOf[KeepDims]) 1 else 0))
       val allInputs             = Tuple1(data)
       (callOp(name, "ArgMin", allInputs, map))
     }