feat: add onnx ArgMin node

Author: tye-singwa

crates/burn-import/onnx-tests/build.rs

@@ -10,6 +10,7 @@ fn main() {
         .input("tests/and/and.onnx")
         .input("tests/add/add_int.onnx")
         .input("tests/argmax/argmax.onnx")
+        .input("tests/argmin/argmin.onnx")
         .input("tests/avg_pool1d/avg_pool1d.onnx")
         .input("tests/avg_pool2d/avg_pool2d.onnx")
         .input("tests/batch_norm/batch_norm.onnx")

crates/burn-import/onnx-tests/tests/argmin/argmin.onnx

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+
+# used to generate model: onnx-tests/tests/argmin/argmin.onnx
+
+import torch
+import torch.nn as nn
+
+
+class Model(nn.Module):
+    def __init__(self, argmin_dim: int = 0):
+        super(Model, self).__init__()
+        self._argmin_dim = argmin_dim
+
+    def forward(self, x):
+        # Note: only keepdim=True is supported in burn
+        y = torch.argmin(input=x, dim=self._argmin_dim, keepdim=True)
+        return y
+
+
+def main():
+
+    # Export to onnx
+    model = Model(1)
+    model.eval()
+    device = torch.device("cpu")
+    onnx_name = "argmin.onnx"
+    dummy_input = torch.randn((3, 4), device=device)
+    torch.onnx.export(model, dummy_input, onnx_name,
+                      verbose=False, opset_version=16)
+
+    print("Finished exporting model to {}".format(onnx_name))
+
+    # Output some test data for use in the test
+    test_input = torch.randn((2, 3), device=device)
+    print("Test input data shape: {}".format(test_input.shape))
+    print("Test input data:\n{}".format(test_input))
+    output = model.forward(test_input)
+
+    print("Test output data shape: {}".format(output.shape))
+    print("Test output data:\n{}".format(output))
+
+
+if __name__ == '__main__':
+    main()

crates/burn-import/onnx-tests/tests/argmin/argmin.py

@@ -0,0 +1,28 @@
+// Import the shared macro
+use crate::include_models;
+include_models!(argmin);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use burn::tensor::{Tensor, TensorData};
+
+    type Backend = burn_ndarray::NdArray<f32>;
+
+    #[test]
+    fn argmin() {
+        // Initialize the model with weights (loaded from the exported file)
+        let model: argmin::Model<Backend> = argmin::Model::default();
+
+        let device = Default::default();
+        // Run the model
+        let input = Tensor::<Backend, 2>::from_floats(
+            [[1.6124, 1.0463, -1.3808], [-0.3852, 0.1301, 0.9780]],
+            &device,
+        );
+        let output = model.forward(input);
+        let expected = TensorData::from([[2i64], [0]]);
+
+        output.to_data().assert_eq(&expected, true);
+    }
+}

crates/burn-import/onnx-tests/tests/argmin/mod.rs

@@ -6,6 +6,7 @@ extern crate alloc;
 pub mod add;
 pub mod and;
 pub mod argmax;
+pub mod argmin;
 pub mod avg_pool;
 pub mod batch_norm;
 pub mod cast;

crates/burn-import/onnx-tests/tests/test_mod.rs

@@ -0,0 +1,102 @@
+use super::{Node, NodeCodegen};
+use crate::burn::{TensorKind, TensorType, ToTokens, Type};
+
+use burn::record::PrecisionSettings;
+use quote::quote;
+
+#[derive(Debug, Clone, new)]
+pub struct ArgMinNode {
+    pub input: TensorType,
+    pub output: TensorType,
+    pub axis: usize,
+}
+
+impl<PS: PrecisionSettings> NodeCodegen<PS> for ArgMinNode {
+    fn output_types(&self) -> Vec<Type> {
+        let mut output = self.output.clone();
+        output.kind = TensorKind::Int;
+        vec![Type::Tensor(output)]
+    }
+
+    fn input_types(&self) -> Vec<crate::burn::Type> {
+        vec![Type::Tensor(self.input.clone())]
+    }
+
+    fn forward(
+        &self,
+        scope: &mut crate::burn::Scope,
+        node_position: usize,
+    ) -> proc_macro2::TokenStream {
+        //NOTE: select_last_index and keep_dims are not supported
+        let axis = self.axis.to_tokens();
+
+        let input = scope.tensor_use_owned(&self.input, node_position);
+        let output = &self.output.name;
+
+        quote! {
+            let #output = #input.argmin(#axis);
+        }
+    }
+
+    fn into_node(self) -> super::Node<PS> {
+        Node::ArgMin(self)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use burn::record::FullPrecisionSettings;
+
+    use super::*;
+    use crate::burn::{TensorType, graph::BurnGraph, node::test::assert_tokens};
+
+    #[test]
+    fn test_codegen_argmin() {
+        let mut graph = BurnGraph::<FullPrecisionSettings>::default();
+
+        graph.register(ArgMinNode::new(
+            TensorType::new_float("tensor1", 2),
+            TensorType::new_int("tensor2", 2),
+            1,
+        ));
+
+        graph.register_input_output(vec!["tensor1".to_string()], vec!["tensor2".to_string()]);
+
+        let expected = quote! {
+            use burn::tensor::Int;
+            use burn::{
+                module::Module,
+                tensor::{backend::Backend, Tensor},
+            };
+
+            #[derive(Module, Debug)]
+            pub struct Model<B: Backend> {
+                phantom: core::marker::PhantomData<B>,
+                device: burn::module::Ignored<B::Device>,
+            }
+
+            impl<B: Backend> Model <B> {
+                #[allow(unused_variables)]
+                pub fn new(device: &B::Device) -> Self {
+                    Self {
+                        phantom: core::marker::PhantomData,
+                        device: burn::module::Ignored(device.clone()),
+                    }
+                }
+
+                #[allow(clippy::let_and_return, clippy::approx_constant)]
+                pub fn forward(
+                    &self,
+                    tensor1: Tensor<B, 2>
+                ) -> Tensor<B, 2, Int> {
+                    let tensor2 = tensor1.argmin(1);
+
+                    tensor2
+                }
+            }
+        };
+
+        assert_tokens(graph.codegen(), expected);
+    }
+}

crates/burn-import/src/burn/node/argmin.rs

@@ -1,7 +1,7 @@
 use std::marker::PhantomData;
 
 use super::{
-    argmax::ArgMaxNode, avg_pool1d::AvgPool1dNode, avg_pool2d::AvgPool2dNode,
+    argmax::ArgMaxNode, argmin::ArgMinNode, avg_pool1d::AvgPool1dNode, avg_pool2d::AvgPool2dNode,
     batch_norm::BatchNormNode, binary::BinaryNode, ceil::CeilNode, clip::ClipNode,
     concat::ConcatNode, constant::ConstantNode, constant_of_shape::ConstantOfShapeNode,
     conv_transpose_1d::ConvTranspose1dNode, conv_transpose_2d::ConvTranspose2dNode,
@@ -86,6 +86,7 @@ pub trait NodeCodegen<PS: PrecisionSettings>: std::fmt::Debug {
 #[derive(Debug, Clone)]
 pub enum Node<PS: PrecisionSettings> {
     ArgMax(ArgMaxNode),
+    ArgMin(ArgMinNode),
     AvgPool1d(AvgPool1dNode),
     AvgPool2d(AvgPool2dNode),
     BatchNorm(BatchNormNode),
@@ -147,6 +148,7 @@ macro_rules! match_all {
         #[allow(clippy::redundant_closure_call)]
         match $self {
             Node::ArgMax(node) => $func(node),
+            Node::ArgMin(node) => $func(node),
             Node::AvgPool1d(node) => $func(node),
             Node::AvgPool2d(node) => $func(node),
             Node::BatchNorm(node) => $func(node),
@@ -216,6 +218,7 @@ impl<PS: PrecisionSettings> Node<PS> {
     pub fn name(&self) -> &str {
         match self {
             Node::ArgMax(_) => "argmax",
+            Node::ArgMin(_) => "argmin",
             Node::AvgPool1d(_) => "avg_pool1d",
             Node::AvgPool2d(_) => "avg_pool2d",
             Node::BatchNorm(_) => "batch_norm",

crates/burn-import/src/burn/node/base.rs

@@ -1,6 +1,7 @@
 mod base;
 
 pub(crate) mod argmax;
+pub(crate) mod argmin;
 pub(crate) mod avg_pool1d;
 pub(crate) mod avg_pool2d;
 pub(crate) mod batch_norm;

crates/burn-import/src/burn/node/mod.rs

@@ -17,6 +17,7 @@ use crate::{
         graph::BurnGraph,
         node::{
             argmax::ArgMaxNode,
+            argmin::ArgMinNode,
             avg_pool1d::AvgPool1dNode,
             avg_pool2d::AvgPool2dNode,
             batch_norm::BatchNormNode,
@@ -80,17 +81,18 @@ use onnx_ir::{
         TensorType as OnnxTensorType,
     },
     node::{
-        argmax::argmax_config, avg_pool1d::avg_pool1d_config, avg_pool2d::avg_pool2d_config,
-        batch_norm::batch_norm_config, clip::clip_config, concat::concat_config,
-        conv_transpose1d::conv_transpose1d_config, conv_transpose2d::conv_transpose2d_config,
-        conv_transpose3d::conv_transpose3d_config, conv1d::conv1d_config, conv2d::conv2d_config,
-        conv3d::conv3d_config, dropout::dropout_config, expand::expand_config,
-        flatten::flatten_config, gather::gather_config, gemm::gemm_config,
-        group_norm::group_norm_config, hard_sigmoid::hard_sigmoid_config,
-        instance_norm::instance_norm_config, layer_norm::layer_norm_config,
-        leaky_relu::leaky_relu_config, linear::linear_config, log_softmax::log_softmax_config,
-        max_pool1d::max_pool1d_config, max_pool2d::max_pool2d_config, one_hot::one_hot_config,
-        pad::pad_config, reduce_max::reduce_max_config, reduce_mean::reduce_mean_config,
+        argmax::argmax_config, argmin::argmin_config, avg_pool1d::avg_pool1d_config,
+        avg_pool2d::avg_pool2d_config, batch_norm::batch_norm_config, clip::clip_config,
+        concat::concat_config, conv_transpose1d::conv_transpose1d_config,
+        conv_transpose2d::conv_transpose2d_config, conv_transpose3d::conv_transpose3d_config,
+        conv1d::conv1d_config, conv2d::conv2d_config, conv3d::conv3d_config,
+        dropout::dropout_config, expand::expand_config, flatten::flatten_config,
+        gather::gather_config, gemm::gemm_config, group_norm::group_norm_config,
+        hard_sigmoid::hard_sigmoid_config, instance_norm::instance_norm_config,
+        layer_norm::layer_norm_config, leaky_relu::leaky_relu_config, linear::linear_config,
+        log_softmax::log_softmax_config, max_pool1d::max_pool1d_config,
+        max_pool2d::max_pool2d_config, one_hot::one_hot_config, pad::pad_config,
+        reduce_max::reduce_max_config, reduce_mean::reduce_mean_config,
         reduce_min::reduce_min_config, reduce_prod::reduce_prod_config,
         reduce_sum::reduce_sum_config, reshape::reshape_config, resize::resize_config,
         slice::slice_config, softmax::softmax_config, split::split_config, squeeze::squeeze_config,
@@ -287,6 +289,7 @@ impl ParsedOnnxGraph {
             match node.node_type {
                 NodeType::Add => graph.register(Self::add_conversion(node)),
                 NodeType::ArgMax => graph.register(Self::argmax_conversion(node)),
+                NodeType::ArgMin => graph.register(Self::argmin_conversion(node)),
                 NodeType::Sub => graph.register(Self::sub_conversion(node)),
                 NodeType::Mul => graph.register(Self::mul_conversion(node)),
                 NodeType::Div => graph.register(Self::div_conversion(node)),
@@ -968,6 +971,14 @@ impl ParsedOnnxGraph {
         ArgMaxNode::new(input, output, axis)
     }
 
+    fn argmin_conversion(node: Node) -> ArgMinNode {
+        let input = TensorType::from(node.inputs.first().unwrap());
+        let output = TensorType::from(node.outputs.first().unwrap());
+        let axis = argmin_config(&node);
+
+        ArgMinNode::new(input, output, axis)
+    }
+
     fn concat_conversion(node: Node) -> ConcatNode {
         let inputs = node.inputs.iter().map(TensorType::from).collect();