diff --git a/docs/framework/operators/tensor/tensor.argmax.md b/docs/framework/operators/tensor/tensor.argmax.md
index cb3457a45..43b1e3721 100644
--- a/docs/framework/operators/tensor/tensor.argmax.md
+++ b/docs/framework/operators/tensor/tensor.argmax.md
@@ -1,7 +1,7 @@
 # tensor.argmax
 
 ```rust 
-   fn argmax(self: @Tensor<T>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>) -> Tensor<usize>;
+   fn argmax(self: @Tensor<T>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>) -> Tensor<i32>;
 ```
 
 Returns the index of the maximum value along the specified axis.
@@ -9,7 +9,7 @@ Returns the index of the maximum value along the specified axis.
 ## Args
 
 * `self`(`@Tensor<T>`) - The input tensor.
-* `axis`(`usize`) - The axis along which to compute the argmax.
+* `axis`(`i32`) - The axis along which to compute the argmax.
 * `keepdims`(`Option<bool>`) - If true, retains reduced dimensions with length 1. Defaults to true.
 * `select_last_index`(`Option<bool>`) - If true, the index of the last occurrence of the maximum value is returned. Defaults to false.   
 
diff --git a/nodegen/node/argmax.py b/nodegen/node/argmax.py
index c874b2eea..69132adfe 100644
--- a/nodegen/node/argmax.py
+++ b/nodegen/node/argmax.py
@@ -3,673 +3,134 @@
 from ..helpers import make_test, to_fp, Tensor, Dtype, FixedImpl
 
 
-def argmax_use_numpy(data: np.ndarray, axis: int = 0, keepdims: int = 1, dtype=np.int64) -> np.ndarray:
+def argmax_use_numpy(data: np.ndarray, axis: int = 0, keepdims: int = 1) -> np.ndarray:
     result = np.argmax(data, axis=axis)
     if keepdims == 1:
         result = np.expand_dims(result, axis)
-    return result.astype(dtype)
+    return result.astype(np.int64)
 
 
 def argmax_use_numpy_select_last_index(
-    data: np.ndarray, axis: int = 0, keepdims: int = True, dtype=np.int64
+    data: np.ndarray, axis: int = 0, keepdims: int = True
 ) -> np.ndarray:
     data = np.flip(data, axis)
     result = np.argmax(data, axis=axis)
     result = data.shape[axis] - result - 1
     if keepdims:
         result = np.expand_dims(result, axis)
-    return result.astype(dtype)
+    return result.astype(np.int64)
 
 
 class Argmax(RunAll):
 
     @staticmethod
-    def argmax_u32():
-        def argmax_1D():
-            def default_params():
-                x = np.random.randint(0, 255, (3)).astype(np.uint32)
-                y = argmax_use_numpy(x, dtype=np.uint32).reshape((1))
+    def no_keepdims():
+        data = np.array([[2, 1], [3, 10]], dtype=np.float32)
+        axis = 1
+        keepdims = 0
+        result = argmax_use_numpy(data, axis=axis, keepdims=keepdims)
 
-                x = Tensor(Dtype.U32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
+        x = Tensor(Dtype.FP16x16, data.shape, data.flatten())
+        y = Tensor(Dtype.I32, result.shape, result.flatten())
 
-                name = "argmax_u32_1D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = np.random.randint(0, 255, (3)).astype(np.uint32)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.U32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_u32_1D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = np.random.randint(0, 255, (3)).astype(np.uint32)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.U32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_u32_1D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
-
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_1D()
-
-        def argmax_2D():
-            def default_params():
-                x = np.random.randint(0, 255, (2, 2)).astype(np.uint32)
-                y = argmax_use_numpy(x, dtype=np.uint32)
-
-                x = Tensor(Dtype.U32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_u32_2D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = np.random.randint(0, 255, (2, 2)).astype(np.uint32)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32)
-
-                x = Tensor(Dtype.U32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_u32_2D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = np.random.randint(0, 255, (2, 2)).astype(np.uint32)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.uint32)
-
-                x = Tensor(Dtype.U32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_u32_2D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
-
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_2D()
-
-        def argmax_3D():
-            def default_params():
-                x = np.random.randint(0, 255, (2, 2, 2)).astype(np.uint32)
-                y = argmax_use_numpy(x, dtype=np.uint32)
-
-                x = Tensor(Dtype.U32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_u32_3D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = np.random.randint(0, 255, (2, 2, 2)).astype(np.uint32)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32)
-
-                x = Tensor(Dtype.U32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_u32_3D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = np.random.randint(0, 255, (2, 2, 2)).astype(np.uint32)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.uint32)
-
-                x = Tensor(Dtype.U32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_u32_3D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
-
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_3D()
+        name = "argmax_no_keepdims"
+        make_test(
+            [x], y, "input_0.argmax(1, Option::Some(false), Option::None(()))", name)
 
     @staticmethod
-    def argmax_i32():
-        def argmax_1D():
-            def default_params():
-                x = np.random.randint(-127, 127, (3)).astype(np.int32)
-                y = argmax_use_numpy(x, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.I32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i32_1D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = np.random.randint(-127, 127, (3)).astype(np.int32)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.I32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i32_1D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = np.random.randint(0, 255, (3)).astype(np.int32)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.uint32).reshape((1))
+    def keepdims():
+        data = np.array([[2, 1], [3, 10]], dtype=np.float32)
+        axis = 1
+        keepdims = 1
+        result = argmax_use_numpy(data, axis=axis, keepdims=keepdims)
 
-                x = Tensor(Dtype.I32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
+        x = Tensor(Dtype.FP16x16, data.shape, data.flatten())
+        y = Tensor(Dtype.I32, result.shape, result.flatten())
 
-                name = "argmax_i32_1D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
-
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_1D()
-
-        def argmax_2D():
-            def default_params():
-                x = np.random.randint(-127, 127, (2, 2)).astype(np.int32)
-                y = argmax_use_numpy(x, dtype=np.uint32)
-
-                x = Tensor(Dtype.I32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i32_2D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = np.random.randint(-127, 127, (2, 2)).astype(np.int32)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32)
-
-                x = Tensor(Dtype.I32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i32_2D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = np.random.randint(-127, 127, (2, 2)).astype(np.int32)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.int32)
-
-                x = Tensor(Dtype.I32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i32_2D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
-
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_2D()
-
-        def argmax_3D():
-            def default_params():
-                x = np.random.randint(-127, 127, (2, 2, 2)).astype(np.int32)
-                y = argmax_use_numpy(x, dtype=np.uint32)
-
-                x = Tensor(Dtype.I32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i32_3D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = np.random.randint(-127, 127, (2, 2, 2)).astype(np.int32)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32)
-
-                x = Tensor(Dtype.I32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i32_3D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = np.random.randint(-127, 127, (2, 2, 2)).astype(np.int32)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.uint32)
-
-                x = Tensor(Dtype.I32, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i32_3D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
-
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_3D()
+        name = "argmax_keepdims"
+        make_test(
+            [x], y, "input_0.argmax(1, Option::Some(true), Option::None(()))", name)
 
     @staticmethod
-    def argmax_i8():
-        def argmax_1D():
-            def default_params():
-                x = np.random.randint(-127, 127, (3)).astype(np.int8)
-                y = argmax_use_numpy(x, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.I8, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i8_1D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = np.random.randint(-127, 127, (3)).astype(np.int8)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.I8, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i8_1D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = np.random.randint(0, 255, (3)).astype(np.int8)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.I8, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i8_1D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
-
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_1D()
-
-        def argmax_2D():
-            def default_params():
-                x = np.random.randint(-127, 127, (2, 2)).astype(np.int8)
-                y = argmax_use_numpy(x, dtype=np.uint32)
-
-                x = Tensor(Dtype.I8, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i8_2D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
+    def default_axes_keepdims():
+        data = np.array([[2, 1], [3, 10]], dtype=np.float32)
+        keepdims = 1
+        result = argmax_use_numpy(data, keepdims=keepdims)
 
-            def keepdims_false():
-                x = np.random.randint(-127, 127, (2, 2)).astype(np.int8)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32)
+        x = Tensor(Dtype.FP16x16, data.shape, data.flatten())
+        y = Tensor(Dtype.I32, result.shape, result.flatten())
 
-                x = Tensor(Dtype.I8, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i8_2D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = np.random.randint(-127, 127, (2, 2)).astype(np.int8)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.int8)
-
-                x = Tensor(Dtype.I8, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i8_2D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
-
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_2D()
-
-        def argmax_3D():
-            def default_params():
-                x = np.random.randint(-127, 127, (2, 2, 2)).astype(np.int8)
-                y = argmax_use_numpy(x, dtype=np.uint32)
-
-                x = Tensor(Dtype.I8, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i8_3D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = np.random.randint(-127, 127, (2, 2, 2)).astype(np.int8)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32)
-
-                x = Tensor(Dtype.I8, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i8_3D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = np.random.randint(-127, 127, (2, 2, 2)).astype(np.int8)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.uint32)
-
-                x = Tensor(Dtype.I8, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_i8_3D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
-
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_3D()
+        name = "argmax_default_axes_keepdims"
+        make_test(
+            [x], y, "input_0.argmax(0, Option::Some(true), Option::None(()))", name)
 
     @staticmethod
-    def argmax_fp16x16():
-        def argmax_1D():
-            def default_params():
-                x = to_fp(np.random.randint(-127, 127, (3)
-                                            ).astype(np.int8), FixedImpl.FP16x16)
-                y = argmax_use_numpy(x, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.FP16x16, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp16x16_1D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = to_fp(np.random.randint(-127, 127, (3)
-                                            ).astype(np.int8), FixedImpl.FP16x16)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.FP16x16, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp16x16_1D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = to_fp(np.random.randint(0, 255, (3)).astype(
-                    np.int8), FixedImpl.FP16x16)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.FP16x16, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp16x16_1D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
-
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_1D()
-
-        def argmax_2D():
-            def default_params():
-                x = to_fp(np.random.randint(-127, 127, (2, 2)
-                                            ).astype(np.int8), FixedImpl.FP16x16)
-                y = argmax_use_numpy(x, dtype=np.uint32)
-
-                x = Tensor(Dtype.FP16x16, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp16x16_2D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = to_fp(np.random.randint(-127, 127, (2, 2)
-                                            ).astype(np.int8), FixedImpl.FP16x16)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32)
-
-                x = Tensor(Dtype.FP16x16, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp16x16_2D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = to_fp(np.random.randint(-127, 127, (2, 2)
-                                            ).astype(np.int8), FixedImpl.FP16x16)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.int8)
-
-                x = Tensor(Dtype.FP16x16, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp16x16_2D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
+    def negative_axis_keepdims():
+        data = np.array([[2, 1], [3, 10]], dtype=np.float32)
+        axis = -1
+        keepdims = 1
+        result = argmax_use_numpy(data, axis=axis, keepdims=keepdims)
 
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_2D()
+        x = Tensor(Dtype.FP16x16, data.shape, data.flatten())
+        y = Tensor(Dtype.I32, result.shape, result.flatten())
 
-        def argmax_3D():
-            def default_params():
-                x = to_fp(np.random.randint(-127, 127, (2, 2, 2)
-                                            ).astype(np.int8), FixedImpl.FP16x16)
-                y = argmax_use_numpy(x, dtype=np.uint32)
-
-                x = Tensor(Dtype.FP16x16, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp16x16_3D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = to_fp(np.random.randint(-127, 127, (2, 2, 2)
-                                            ).astype(np.int8), FixedImpl.FP16x16)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32)
-
-                x = Tensor(Dtype.FP16x16, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp16x16_3D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = to_fp(np.random.randint(-127, 127, (2, 2, 2)
-                                            ).astype(np.int8), FixedImpl.FP16x16)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.uint32)
-
-                x = Tensor(Dtype.FP16x16, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp16x16_3D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
-
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_3D()
+        name = "argmax_negative_axis_keepdims"
+        make_test(
+            [x], y, "input_0.argmax(-1, Option::Some(true), Option::None(()))", name)
 
     @staticmethod
-    def argmax_fp8x23():
-        def argmax_1D():
-            def default_params():
-                x = to_fp(np.random.randint(-127, 127, (3)
-                                            ).astype(np.int8), FixedImpl.FP8x23)
-                y = argmax_use_numpy(x, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.FP8x23, x.shape,
-                           x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp8x23_1D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = to_fp(np.random.randint(-127, 127, (3)
-                                            ).astype(np.int8), FixedImpl.FP8x23)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.FP8x23, x.shape,
-                           x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp8x23_1D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = to_fp(np.random.randint(0, 255, (3)).astype(
-                    np.int8), FixedImpl.FP8x23)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.uint32).reshape((1))
-
-                x = Tensor(Dtype.FP8x23, x.shape,
-                           x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
+    def no_keepdims_select_last_index():
+        data = np.array([[2, 2], [3, 10]], dtype=np.float32)
+        axis = 1
+        keepdims = 0
+        result = argmax_use_numpy_select_last_index(
+            data, axis=axis, keepdims=keepdims)
 
-                name = "argmax_fp8x23_1D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
+        x = Tensor(Dtype.FP16x16, data.shape, data.flatten())
+        y = Tensor(Dtype.I32, result.shape, result.flatten())
 
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_1D()
+        name = "argmax_no_keepdims_select_last_index"
+        make_test(
+            [x], y, "input_0.argmax(1, Option::Some(false), Option::Some(true))", name)
 
-        def argmax_2D():
-            def default_params():
-                x = to_fp(np.random.randint(-127, 127, (2, 2)
-                                            ).astype(np.int8), FixedImpl.FP8x23)
-                y = argmax_use_numpy(x, dtype=np.uint32)
-
-                x = Tensor(Dtype.FP8x23, x.shape, x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp8x23_2D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = to_fp(np.random.randint(-127, 127, (2, 2)
-                                            ).astype(np.int8), FixedImpl.FP8x23)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32)
-
-                x = Tensor(Dtype.FP8x23, x.shape,
-                           x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp8x23_2D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
-
-            def last_index():
-                x = to_fp(np.random.randint(-127, 127, (2, 2)
-                                            ).astype(np.int8), FixedImpl.FP8x23)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.int8)
-
-                x = Tensor(Dtype.FP8x23, x.shape,
-                           x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp8x23_2D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
-
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_2D()
-
-        def argmax_3D():
-            def default_params():
-                x = to_fp(np.random.randint(-127, 127, (2, 2, 2)
-                                            ).astype(np.int8), FixedImpl.FP8x23)
-                y = argmax_use_numpy(x, dtype=np.uint32)
-
-                x = Tensor(Dtype.FP8x23, x.shape,
-                           x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
-
-                name = "argmax_fp8x23_3D_default"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::None(()))", name)
-
-            def keepdims_false():
-                x = to_fp(np.random.randint(-127, 127, (2, 2, 2)
-                                            ).astype(np.int8), FixedImpl.FP8x23)
-                y = argmax_use_numpy(
-                    x, keepdims=0, dtype=np.uint32)
+    @staticmethod
+    def keepdims_select_last_index():
+        data = np.array([[2, 2], [3, 10]], dtype=np.float32)
+        axis = 1
+        keepdims = 1
+        result = argmax_use_numpy_select_last_index(
+            data, axis=axis, keepdims=keepdims)
 
-                x = Tensor(Dtype.FP8x23, x.shape,
-                           x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
+        x = Tensor(Dtype.FP16x16, data.shape, data.flatten())
+        y = Tensor(Dtype.I32, result.shape, result.flatten())
 
-                name = "argmax_fp8x23_3D_keepdims_false"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::Some(false), Option::None(()))", name)
+        name = "argmax_keepdims_select_last_index"
+        make_test(
+            [x], y, "input_0.argmax(1, Option::Some(true), Option::Some(true))", name)
 
-            def last_index():
-                x = to_fp(np.random.randint(-127, 127, (2, 2, 2)
-                                            ).astype(np.int8), FixedImpl.FP8x23)
-                y = argmax_use_numpy_select_last_index(
-                    x, dtype=np.uint32)
+    @staticmethod
+    def default_axes_keepdims_select_last_index():
+        data = np.array([[2, 2], [3, 10]], dtype=np.float32)
+        keepdims = 1
+        result = argmax_use_numpy_select_last_index(data, keepdims=keepdims)
 
-                x = Tensor(Dtype.FP8x23, x.shape,
-                           x.flatten())
-                y = Tensor(Dtype.U32, y.shape, y.flatten())
+        x = Tensor(Dtype.FP16x16, data.shape, data.flatten())
+        y = Tensor(Dtype.I32, result.shape, result.flatten())
 
-                name = "argmax_fp8x23_3D_last_index"
-                make_test(
-                    [x], y, "input_0.argmax(0, Option::None(()), Option::Some(true))", name)
+        name = "argmax_default_axes_keepdims_select_last_index"
+        make_test(
+            [x], y, "input_0.argmax(0, Option::Some(true), Option::Some(true))", name)
 
-            default_params()
-            keepdims_false()
-            last_index()
-        argmax_3D()
+    @staticmethod
+    def negative_axis_keepdims_select_last_index():
+        data = np.array([[2, 2], [3, 10]], dtype=np.float32)
+        axis = -1
+        keepdims = 1
+        result = argmax_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims)
+
+        x = Tensor(Dtype.FP16x16, data.shape, data.flatten())
+        y = Tensor(Dtype.I32, result.shape, result.flatten())
+
+        name = "argmax_negative_axis_keepdims_select_last_index"
+        make_test(
+            [x], y, "input_0.argmax(-1, Option::Some(true), Option::Some(true))", name)
diff --git a/src/operators/ml/linear/linear_classifier.cairo b/src/operators/ml/linear/linear_classifier.cairo
index 18cbd6988..003cf742b 100644
--- a/src/operators/ml/linear/linear_classifier.cairo
+++ b/src/operators/ml/linear/linear_classifier.cairo
@@ -197,7 +197,9 @@ impl LinearClassifierImpl<
             let mut labels = scores.argmax(1, Option::None, Option::None);
             loop {
                 match labels.data.pop_front() {
-                    Option::Some(i) => { labels_list.append(*classlabels[*i]); },
+                    Option::Some(i) => {
+                        labels_list.append(*classlabels[(*i).try_into().unwrap()]);
+                    },
                     Option::None => { break; }
                 };
             };
@@ -205,48 +207,56 @@ impl LinearClassifierImpl<
             let mut i = 0;
             match classifier.post_transform {
                 POST_TRANSFORM::NONE => {
-                    while i != scores.data.len() {
-                        if *scores.data.at(i) >= NumberTrait::zero() {
-                            labels_list.append(*classlabels[0]);
-                        } else {
-                            labels_list.append(0);
-                        }
+                    while i != scores
+                        .data
+                        .len() {
+                            if *scores.data.at(i) >= NumberTrait::zero() {
+                                labels_list.append(*classlabels[0]);
+                            } else {
+                                labels_list.append(0);
+                            }
 
-                        i += 1;
-                    };
+                            i += 1;
+                        };
                 },
                 POST_TRANSFORM::SOFTMAX => {
-                    while i != scores.data.len() {
-                        if *scores.data.at(i) >= NumberTrait::half() {
-                            labels_list.append(*classlabels[0]);
-                        } else {
-                            labels_list.append(0);
-                        }
+                    while i != scores
+                        .data
+                        .len() {
+                            if *scores.data.at(i) >= NumberTrait::half() {
+                                labels_list.append(*classlabels[0]);
+                            } else {
+                                labels_list.append(0);
+                            }
 
-                        i += 1;
-                    };
+                            i += 1;
+                        };
                 },
                 POST_TRANSFORM::LOGISTIC => {
-                    while i != scores.data.len() {
-                        if *scores.data.at(i) >= NumberTrait::half() {
-                            labels_list.append(*classlabels[0]);
-                        } else {
-                            labels_list.append(0);
-                        }
+                    while i != scores
+                        .data
+                        .len() {
+                            if *scores.data.at(i) >= NumberTrait::half() {
+                                labels_list.append(*classlabels[0]);
+                            } else {
+                                labels_list.append(0);
+                            }
 
-                        i += 1;
-                    };
+                            i += 1;
+                        };
                 },
                 POST_TRANSFORM::SOFTMAXZERO => {
-                    while i != scores.data.len() {
-                        if *scores.data.at(i) >= NumberTrait::half() {
-                            labels_list.append(*classlabels[0]);
-                        } else {
-                            labels_list.append(0);
-                        }
+                    while i != scores
+                        .data
+                        .len() {
+                            if *scores.data.at(i) >= NumberTrait::half() {
+                                labels_list.append(*classlabels[0]);
+                            } else {
+                                labels_list.append(0);
+                            }
 
-                        i += 1;
-                    };
+                            i += 1;
+                        };
                 },
                 POST_TRANSFORM::PROBIT => core::panic_with_felt252('Probit not supported yet'),
             };
diff --git a/src/operators/tensor/core.cairo b/src/operators/tensor/core.cairo
index 70f551ec8..adb99299e 100644
--- a/src/operators/tensor/core.cairo
+++ b/src/operators/tensor/core.cairo
@@ -687,7 +687,7 @@ trait TensorTrait<T> {
     /// # tensor.argmax
     ///
     /// ```rust 
-    ///    fn argmax(self: @Tensor<T>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>) -> Tensor<usize>;
+    ///    fn argmax(self: @Tensor<T>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>) -> Tensor<i32>;
     /// ```
     ///
     /// Returns the index of the maximum value along the specified axis.
@@ -695,7 +695,7 @@ trait TensorTrait<T> {
     /// ## Args
     ///
     /// * `self`(`@Tensor<T>`) - The input tensor.
-    /// * `axis`(`usize`) - The axis along which to compute the argmax.
+    /// * `axis`(`i32`) - The axis along which to compute the argmax.
     /// * `keepdims`(`Option<bool>`) - If true, retains reduced dimensions with length 1. Defaults to true.
     /// * `select_last_index`(`Option<bool>`) - If true, the index of the last occurrence of the maximum value is returned. Defaults to false.   
     ///
@@ -765,8 +765,8 @@ trait TensorTrait<T> {
     /// ```
     ///
     fn argmax(
-        self: @Tensor<T>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>
-    ) -> Tensor<usize>;
+        self: @Tensor<T>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>
+    ) -> Tensor<i32>;
     /// # tensor.argmin
     ///
     /// ```rust 
diff --git a/src/operators/tensor/implementations/tensor_bool.cairo b/src/operators/tensor/implementations/tensor_bool.cairo
index a6c54261c..d08a0a2c1 100644
--- a/src/operators/tensor/implementations/tensor_bool.cairo
+++ b/src/operators/tensor/implementations/tensor_bool.cairo
@@ -78,8 +78,8 @@ impl BoolTensor of TensorTrait<bool> {
     }
 
     fn argmax(
-        self: @Tensor<bool>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>
-    ) -> Tensor<usize> {
+        self: @Tensor<bool>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>
+    ) -> Tensor<i32> {
         panic(array!['not supported!'])
     }
 
diff --git a/src/operators/tensor/implementations/tensor_complex64.cairo b/src/operators/tensor/implementations/tensor_complex64.cairo
index fa465b33e..2bbbab8fe 100644
--- a/src/operators/tensor/implementations/tensor_complex64.cairo
+++ b/src/operators/tensor/implementations/tensor_complex64.cairo
@@ -90,10 +90,10 @@ impl Complex64Tensor of TensorTrait<complex64> {
 
     fn argmax(
         self: @Tensor<complex64>,
-        axis: usize,
+        axis: i32,
         keepdims: Option<bool>,
         select_last_index: Option<bool>
-    ) -> Tensor<usize> {
+    ) -> Tensor<i32> {
         panic(array!['not supported!'])
     }
 
diff --git a/src/operators/tensor/implementations/tensor_fp16x16.cairo b/src/operators/tensor/implementations/tensor_fp16x16.cairo
index e1795000b..e1a780e17 100644
--- a/src/operators/tensor/implementations/tensor_fp16x16.cairo
+++ b/src/operators/tensor/implementations/tensor_fp16x16.cairo
@@ -89,8 +89,8 @@ impl FP16x16Tensor of TensorTrait<FP16x16> {
     }
 
     fn argmax(
-        self: @Tensor<FP16x16>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>
-    ) -> Tensor<usize> {
+        self: @Tensor<FP16x16>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>
+    ) -> Tensor<i32> {
         math::argmax::argmax(self, axis, keepdims, select_last_index)
     }
 
diff --git a/src/operators/tensor/implementations/tensor_fp16x16wide.cairo b/src/operators/tensor/implementations/tensor_fp16x16wide.cairo
index c2e9b9344..13ea883cf 100644
--- a/src/operators/tensor/implementations/tensor_fp16x16wide.cairo
+++ b/src/operators/tensor/implementations/tensor_fp16x16wide.cairo
@@ -94,10 +94,10 @@ impl FP16x16WTensor of TensorTrait<FP16x16W> {
 
     fn argmax(
         self: @Tensor<FP16x16W>,
-        axis: usize,
+        axis: i32,
         keepdims: Option<bool>,
         select_last_index: Option<bool>
-    ) -> Tensor<usize> {
+    ) -> Tensor<i32> {
         math::argmax::argmax(self, axis, keepdims, select_last_index)
     }
 
diff --git a/src/operators/tensor/implementations/tensor_fp32x32.cairo b/src/operators/tensor/implementations/tensor_fp32x32.cairo
index 5cbf64139..cb745f113 100644
--- a/src/operators/tensor/implementations/tensor_fp32x32.cairo
+++ b/src/operators/tensor/implementations/tensor_fp32x32.cairo
@@ -86,8 +86,8 @@ impl FP32x32Tensor of TensorTrait<FP32x32> {
     }
 
     fn argmax(
-        self: @Tensor<FP32x32>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>
-    ) -> Tensor<usize> {
+        self: @Tensor<FP32x32>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>
+    ) -> Tensor<i32> {
         math::argmax::argmax(self, axis, keepdims, select_last_index)
     }
 
diff --git a/src/operators/tensor/implementations/tensor_fp64x64.cairo b/src/operators/tensor/implementations/tensor_fp64x64.cairo
index 3c18ae38a..0ee0f41d2 100644
--- a/src/operators/tensor/implementations/tensor_fp64x64.cairo
+++ b/src/operators/tensor/implementations/tensor_fp64x64.cairo
@@ -86,8 +86,8 @@ impl FP64x64Tensor of TensorTrait<FP64x64> {
     }
 
     fn argmax(
-        self: @Tensor<FP64x64>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>
-    ) -> Tensor<usize> {
+        self: @Tensor<FP64x64>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>
+    ) -> Tensor<i32> {
         math::argmax::argmax(self, axis, keepdims, select_last_index)
     }
 
diff --git a/src/operators/tensor/implementations/tensor_fp8x23.cairo b/src/operators/tensor/implementations/tensor_fp8x23.cairo
index 33a19c080..c404e4641 100644
--- a/src/operators/tensor/implementations/tensor_fp8x23.cairo
+++ b/src/operators/tensor/implementations/tensor_fp8x23.cairo
@@ -86,8 +86,8 @@ impl FP8x23Tensor of TensorTrait<FP8x23> {
     }
 
     fn argmax(
-        self: @Tensor<FP8x23>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>
-    ) -> Tensor<usize> {
+        self: @Tensor<FP8x23>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>
+    ) -> Tensor<i32> {
         math::argmax::argmax(self, axis, keepdims, select_last_index)
     }
 
diff --git a/src/operators/tensor/implementations/tensor_fp8x23wide.cairo b/src/operators/tensor/implementations/tensor_fp8x23wide.cairo
index 459657891..f415266c4 100644
--- a/src/operators/tensor/implementations/tensor_fp8x23wide.cairo
+++ b/src/operators/tensor/implementations/tensor_fp8x23wide.cairo
@@ -89,8 +89,8 @@ impl FP8x23WTensor of TensorTrait<FP8x23W> {
     }
 
     fn argmax(
-        self: @Tensor<FP8x23W>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>
-    ) -> Tensor<usize> {
+        self: @Tensor<FP8x23W>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>
+    ) -> Tensor<i32> {
         math::argmax::argmax(self, axis, keepdims, select_last_index)
     }
 
diff --git a/src/operators/tensor/implementations/tensor_i32.cairo b/src/operators/tensor/implementations/tensor_i32.cairo
index d4fa1b09c..3c3641310 100644
--- a/src/operators/tensor/implementations/tensor_i32.cairo
+++ b/src/operators/tensor/implementations/tensor_i32.cairo
@@ -87,8 +87,8 @@ impl I32Tensor of TensorTrait<i32> {
     }
 
     fn argmax(
-        self: @Tensor<i32>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>
-    ) -> Tensor<usize> {
+        self: @Tensor<i32>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>
+    ) -> Tensor<i32> {
         math::argmax::argmax(self, axis, keepdims, select_last_index)
     }
 
diff --git a/src/operators/tensor/implementations/tensor_i8.cairo b/src/operators/tensor/implementations/tensor_i8.cairo
index 903d66638..a872a4665 100644
--- a/src/operators/tensor/implementations/tensor_i8.cairo
+++ b/src/operators/tensor/implementations/tensor_i8.cairo
@@ -84,8 +84,8 @@ impl I8Tensor of TensorTrait<i8> {
     }
 
     fn argmax(
-        self: @Tensor<i8>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>
-    ) -> Tensor<usize> {
+        self: @Tensor<i8>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>
+    ) -> Tensor<i32> {
         math::argmax::argmax(self, axis, keepdims, select_last_index)
     }
 
diff --git a/src/operators/tensor/implementations/tensor_u32.cairo b/src/operators/tensor/implementations/tensor_u32.cairo
index 01456aec6..e85316a96 100644
--- a/src/operators/tensor/implementations/tensor_u32.cairo
+++ b/src/operators/tensor/implementations/tensor_u32.cairo
@@ -83,8 +83,8 @@ impl U32Tensor of TensorTrait<u32> {
     }
 
     fn argmax(
-        self: @Tensor<u32>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>
-    ) -> Tensor<usize> {
+        self: @Tensor<u32>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>
+    ) -> Tensor<i32> {
         math::argmax::argmax(self, axis, keepdims, select_last_index)
     }
 
diff --git a/src/operators/tensor/math/argmax.cairo b/src/operators/tensor/math/argmax.cairo
index f16c99b5c..1a48d3d8b 100644
--- a/src/operators/tensor/math/argmax.cairo
+++ b/src/operators/tensor/math/argmax.cairo
@@ -1,4 +1,6 @@
-use orion::operators::tensor::core::{Tensor, TensorTrait, ravel_index, unravel_index};
+use core::option::OptionTrait;
+use core::traits::TryInto;
+use orion::operators::tensor::{core::{Tensor, TensorTrait, ravel_index, unravel_index}, I32Tensor};
 use orion::operators::tensor::helpers::{reduce_output_shape, len_from_shape, combine_indices};
 use orion::numbers::NumberTrait;
 
@@ -6,23 +8,22 @@ use orion::numbers::NumberTrait;
 fn argmax<
     T,
     MAG,
-    impl UsizeTensor: TensorTrait<usize>,
     impl TNumber: NumberTrait<T, MAG>,
     impl TPartialOrd: PartialOrd<T>,
     impl TPartialEq: PartialEq<T>,
     impl TCopy: Copy<T>,
     impl TDrop: Drop<T>,
 >(
-    self: @Tensor<T>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>
-) -> Tensor<usize> {
-    let keepdims = match keepdims {
-        Option::Some(val) => val,
-        Option::None => true,
-    };
-
-    let select_last_index = match select_last_index {
-        Option::Some(val) => val,
-        Option::None => false,
+    self: @Tensor<T>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>
+) -> Tensor<i32> {
+    let keepdims = keepdims.unwrap_or(true);
+    let select_last_index = select_last_index.unwrap_or(false);
+
+    // Convert negative axis to positive
+    let axis = if axis < 0 {
+        ((*self.shape).len().try_into().unwrap() + axis).try_into().unwrap()
+    } else {
+        axis.try_into().unwrap()
     };
 
     assert(axis <= (*self.shape).len(), 'axis out of dimensions');
@@ -31,7 +32,7 @@ fn argmax<
         return find_argmax_1D::<T>(*self, axis, true, select_last_index);
     }
 
-    let mut output_data: Array<u32> = array![];
+    let mut output_data: Array<i32> = array![];
 
     let output_shape = reduce_output_shape(*self.shape, axis, false);
     let output_data_len = len_from_shape(output_shape);
@@ -48,7 +49,7 @@ fn argmax<
         index += 1;
     };
 
-    TensorTrait::<usize>::new(reduce_output_shape(*self.shape, axis, keepdims), output_data.span())
+    TensorTrait::<i32>::new(reduce_output_shape(*self.shape, axis, keepdims), output_data.span())
 }
 
 /// Helper function that finds the index of the maximum value in a flat tensor.
@@ -66,21 +67,20 @@ fn argmax<
 /// * A usize value representing the index of the maximum value along the specified axis.
 fn find_argmax_1D<
     T,
-    impl UsizeTensor: TensorTrait<usize>,
     impl TPartialOrd: PartialOrd<T>,
     impl TPartialEq: PartialEq<T>,
     impl TCopy: Copy<T>,
     impl TDrop: Drop<T>,
 >(
     mut input: Tensor<T>, axis: usize, keepdims: bool, select_last_index: bool
-) -> Tensor<usize> {
-    let mut output_data = ArrayTrait::<usize>::new();
+) -> Tensor<i32> {
+    let mut output_data = ArrayTrait::<i32>::new();
 
     let mut max = match input.data.pop_front() {
         Option::Some(item) => *item,
         Option::None => {
             return TensorTrait::<
-                usize
+                i32
             >::new(reduce_output_shape(input.shape, axis, keepdims), output_data.span());
         }
     };
@@ -108,7 +108,7 @@ fn find_argmax_1D<
     output_data.append(max_index);
 
     return TensorTrait::<
-        usize
+        i32
     >::new(reduce_output_shape(input.shape, axis, keepdims), output_data.span());
 }
 
@@ -143,9 +143,9 @@ fn find_argmax<
     max_value: T,
     argmax: usize,
     select_last_index: bool
-) -> usize {
+) -> i32 {
     if axis_index == *(*input.shape)[axis] {
-        return argmax;
+        return argmax.try_into().unwrap();
     }
 
     let input_indices = combine_indices(output_indices, axis_index, axis);
diff --git a/tests/nodes.cairo b/tests/nodes.cairo
index 19b1a7df9..a8b5f3db6 100644
--- a/tests/nodes.cairo
+++ b/tests/nodes.cairo
@@ -16,51 +16,6 @@
 // mod add_i8_broadcast;
 // mod add_u32;
 // mod add_u32_broadcast;
-// mod argmax_fp16x16_1D_default;
-// mod argmax_fp16x16_1D_keepdims_false;
-// mod argmax_fp16x16_1D_last_index;
-// mod argmax_fp16x16_2D_default;
-// mod argmax_fp16x16_2D_keepdims_false;
-// mod argmax_fp16x16_2D_last_index;
-// mod argmax_fp16x16_3D_default;
-// mod argmax_fp16x16_3D_keepdims_false;
-// mod argmax_fp16x16_3D_last_index;
-// mod argmax_fp8x23_1D_default;
-// mod argmax_fp8x23_1D_keepdims_false;
-// mod argmax_fp8x23_1D_last_index;
-// mod argmax_fp8x23_2D_default;
-// mod argmax_fp8x23_2D_keepdims_false;
-// mod argmax_fp8x23_2D_last_index;
-// mod argmax_fp8x23_3D_default;
-// mod argmax_fp8x23_3D_keepdims_false;
-// mod argmax_fp8x23_3D_last_index;
-// mod argmax_i32_1D_default;
-// mod argmax_i32_1D_keepdims_false;
-// mod argmax_i32_1D_last_index;
-// mod argmax_i32_2D_default;
-// mod argmax_i32_2D_keepdims_false;
-// mod argmax_i32_2D_last_index;
-// mod argmax_i32_3D_default;
-// mod argmax_i32_3D_keepdims_false;
-// mod argmax_i32_3D_last_index;
-// mod argmax_i8_1D_default;
-// mod argmax_i8_1D_keepdims_false;
-// mod argmax_i8_1D_last_index;
-// mod argmax_i8_2D_default;
-// mod argmax_i8_2D_keepdims_false;
-// mod argmax_i8_2D_last_index;
-// mod argmax_i8_3D_default;
-// mod argmax_i8_3D_keepdims_false;
-// mod argmax_i8_3D_last_index;
-// mod argmax_u32_1D_default;
-// mod argmax_u32_1D_keepdims_false;
-// mod argmax_u32_1D_last_index;
-// mod argmax_u32_2D_default;
-// mod argmax_u32_2D_keepdims_false;
-// mod argmax_u32_2D_last_index;
-// mod argmax_u32_3D_default;
-// mod argmax_u32_3D_keepdims_false;
-// mod argmax_u32_3D_last_index;
 // mod argmin_fp16x16_1D_default;
 // mod argmin_fp16x16_1D_keepdims_false;
 // mod argmin_fp16x16_1D_last_index;
@@ -1022,3 +977,11 @@ mod softmax_axis_0;
 mod softmax_axis_1;
 mod softmax_axis_2;
 mod softmax_axis_minus_1;
+mod argmax_default_axes_keepdims;
+mod argmax_default_axes_keepdims_select_last_index;
+mod argmax_keepdims;
+mod argmax_keepdims_select_last_index;
+mod argmax_negative_axis_keepdims;
+mod argmax_negative_axis_keepdims_select_last_index;
+mod argmax_no_keepdims;
+mod argmax_no_keepdims_select_last_index;
diff --git a/tests/nodes/argmax_fp16x16_1D_last_index.cairo b/tests/nodes/argmax_default_axes_keepdims.cairo
similarity index 56%
rename from tests/nodes/argmax_fp16x16_1D_last_index.cairo
rename to tests/nodes/argmax_default_axes_keepdims.cairo
index fad13c40f..ba970e3b1 100644
--- a/tests/nodes/argmax_fp16x16_1D_last_index.cairo
+++ b/tests/nodes/argmax_default_axes_keepdims.cairo
@@ -2,21 +2,21 @@ mod input_0;
 mod output_0;
 
 
+use orion::operators::tensor::{I32Tensor, I32TensorAdd};
 use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
 use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::FP16x16TensorPartialEq;
+use orion::utils::{assert_eq, assert_seq_eq};
+use orion::operators::tensor::I32TensorPartialEq;
 use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
+use orion::operators::tensor::FP16x16TensorPartialEq;
 
 #[test]
 #[available_gas(2000000000)]
-fn test_argmax_fp16x16_1D_last_index() {
+fn test_argmax_default_axes_keepdims() {
     let input_0 = input_0::input_0();
-    let z = output_0::output_0();
+    let z_0 = output_0::output_0();
 
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
+    let y_0 = input_0.argmax(0, Option::Some(true), Option::None(()));
 
-    assert_eq(y, z);
+    assert_eq(y_0, z_0);
 }
diff --git a/tests/nodes/argmax_fp16x16_2D_keepdims_false/input_0.cairo b/tests/nodes/argmax_default_axes_keepdims/input_0.cairo
similarity index 65%
rename from tests/nodes/argmax_fp16x16_2D_keepdims_false/input_0.cairo
rename to tests/nodes/argmax_default_axes_keepdims/input_0.cairo
index 906f29fff..2bd0e9ded 100644
--- a/tests/nodes/argmax_fp16x16_2D_keepdims_false/input_0.cairo
+++ b/tests/nodes/argmax_default_axes_keepdims/input_0.cairo
@@ -9,9 +9,9 @@ fn input_0() -> Tensor<FP16x16> {
     shape.append(2);
 
     let mut data = ArrayTrait::new();
-    data.append(FP16x16 { mag: 4521984, sign: false });
-    data.append(FP16x16 { mag: 7667712, sign: true });
-    data.append(FP16x16 { mag: 393216, sign: true });
-    data.append(FP16x16 { mag: 6356992, sign: true });
+    data.append(FP16x16 { mag: 2, sign: false });
+    data.append(FP16x16 { mag: 1, sign: false });
+    data.append(FP16x16 { mag: 3, sign: false });
+    data.append(FP16x16 { mag: 10, sign: false });
     TensorTrait::new(shape.span(), data.span())
 }
diff --git a/tests/nodes/argmax_i32_2D_keepdims_false/input_0.cairo b/tests/nodes/argmax_default_axes_keepdims/output_0.cairo
similarity index 69%
rename from tests/nodes/argmax_i32_2D_keepdims_false/input_0.cairo
rename to tests/nodes/argmax_default_axes_keepdims/output_0.cairo
index 4b1340a16..83db954c1 100644
--- a/tests/nodes/argmax_i32_2D_keepdims_false/input_0.cairo
+++ b/tests/nodes/argmax_default_axes_keepdims/output_0.cairo
@@ -1,16 +1,15 @@
 use core::array::{ArrayTrait, SpanTrait};
 use orion::operators::tensor::{TensorTrait, Tensor};
 use orion::operators::tensor::{I32Tensor, I32TensorAdd};
+use orion::numbers::NumberTrait;
 
-fn input_0() -> Tensor<i32> {
+fn output_0() -> Tensor<i32> {
     let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
+    shape.append(1);
     shape.append(2);
 
     let mut data = ArrayTrait::new();
-    data.append(15);
-    data.append(57);
-    data.append(-83);
-    data.append(13);
+    data.append(1);
+    data.append(1);
     TensorTrait::new(shape.span(), data.span())
 }
diff --git a/tests/nodes/argmax_default_axes_keepdims_select_last_index.cairo b/tests/nodes/argmax_default_axes_keepdims_select_last_index.cairo
new file mode 100644
index 000000000..868a5f4a2
--- /dev/null
+++ b/tests/nodes/argmax_default_axes_keepdims_select_last_index.cairo
@@ -0,0 +1,22 @@
+mod input_0;
+mod output_0;
+
+
+use orion::operators::tensor::{I32Tensor, I32TensorAdd};
+use core::array::{ArrayTrait, SpanTrait};
+use orion::operators::tensor::{TensorTrait, Tensor};
+use orion::utils::{assert_eq, assert_seq_eq};
+use orion::operators::tensor::I32TensorPartialEq;
+use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
+use orion::operators::tensor::FP16x16TensorPartialEq;
+
+#[test]
+#[available_gas(2000000000)]
+fn test_argmax_default_axes_keepdims_select_last_index() {
+    let input_0 = input_0::input_0();
+    let z_0 = output_0::output_0();
+
+    let y_0 = input_0.argmax(0, Option::Some(true), Option::Some(true));
+
+    assert_eq(y_0, z_0);
+}
diff --git a/tests/nodes/argmax_fp16x16_2D_last_index/input_0.cairo b/tests/nodes/argmax_default_axes_keepdims_select_last_index/input_0.cairo
similarity index 65%
rename from tests/nodes/argmax_fp16x16_2D_last_index/input_0.cairo
rename to tests/nodes/argmax_default_axes_keepdims_select_last_index/input_0.cairo
index 708f630e3..c0673e50d 100644
--- a/tests/nodes/argmax_fp16x16_2D_last_index/input_0.cairo
+++ b/tests/nodes/argmax_default_axes_keepdims_select_last_index/input_0.cairo
@@ -9,9 +9,9 @@ fn input_0() -> Tensor<FP16x16> {
     shape.append(2);
 
     let mut data = ArrayTrait::new();
-    data.append(FP16x16 { mag: 1310720, sign: true });
-    data.append(FP16x16 { mag: 3670016, sign: true });
-    data.append(FP16x16 { mag: 6422528, sign: true });
-    data.append(FP16x16 { mag: 5636096, sign: false });
+    data.append(FP16x16 { mag: 2, sign: false });
+    data.append(FP16x16 { mag: 2, sign: false });
+    data.append(FP16x16 { mag: 3, sign: false });
+    data.append(FP16x16 { mag: 10, sign: false });
     TensorTrait::new(shape.span(), data.span())
 }
diff --git a/tests/nodes/argmax_i32_2D_last_index/input_0.cairo b/tests/nodes/argmax_default_axes_keepdims_select_last_index/output_0.cairo
similarity index 69%
rename from tests/nodes/argmax_i32_2D_last_index/input_0.cairo
rename to tests/nodes/argmax_default_axes_keepdims_select_last_index/output_0.cairo
index f06f44db7..83db954c1 100644
--- a/tests/nodes/argmax_i32_2D_last_index/input_0.cairo
+++ b/tests/nodes/argmax_default_axes_keepdims_select_last_index/output_0.cairo
@@ -1,16 +1,15 @@
 use core::array::{ArrayTrait, SpanTrait};
 use orion::operators::tensor::{TensorTrait, Tensor};
 use orion::operators::tensor::{I32Tensor, I32TensorAdd};
+use orion::numbers::NumberTrait;
 
-fn input_0() -> Tensor<i32> {
+fn output_0() -> Tensor<i32> {
     let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
+    shape.append(1);
     shape.append(2);
 
     let mut data = ArrayTrait::new();
-    data.append(-31);
-    data.append(-9);
-    data.append(-103);
-    data.append(-88);
+    data.append(1);
+    data.append(1);
     TensorTrait::new(shape.span(), data.span())
 }
diff --git a/tests/nodes/argmax_fp16x16_1D_default/input_0.cairo b/tests/nodes/argmax_fp16x16_1D_default/input_0.cairo
deleted file mode 100644
index 1d26625fd..000000000
--- a/tests/nodes/argmax_fp16x16_1D_default/input_0.cairo
+++ /dev/null
@@ -1,15 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
-use orion::numbers::{FixedTrait, FP16x16};
-
-fn input_0() -> Tensor<FP16x16> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP16x16 { mag: 4194304, sign: true });
-    data.append(FP16x16 { mag: 917504, sign: false });
-    data.append(FP16x16 { mag: 5832704, sign: false });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp16x16_1D_default/output_0.cairo b/tests/nodes/argmax_fp16x16_1D_default/output_0.cairo
deleted file mode 100644
index 23846d97b..000000000
--- a/tests/nodes/argmax_fp16x16_1D_default/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(2);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp16x16_1D_keepdims_false/output_0.cairo b/tests/nodes/argmax_fp16x16_1D_keepdims_false/output_0.cairo
deleted file mode 100644
index e4dbc507c..000000000
--- a/tests/nodes/argmax_fp16x16_1D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp16x16_1D_last_index/input_0.cairo b/tests/nodes/argmax_fp16x16_1D_last_index/input_0.cairo
deleted file mode 100644
index 9f3f8fb9f..000000000
--- a/tests/nodes/argmax_fp16x16_1D_last_index/input_0.cairo
+++ /dev/null
@@ -1,15 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
-use orion::numbers::{FixedTrait, FP16x16};
-
-fn input_0() -> Tensor<FP16x16> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP16x16 { mag: 2359296, sign: true });
-    data.append(FP16x16 { mag: 5701632, sign: false });
-    data.append(FP16x16 { mag: 6684672, sign: false });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp16x16_1D_last_index/output_0.cairo b/tests/nodes/argmax_fp16x16_1D_last_index/output_0.cairo
deleted file mode 100644
index 23846d97b..000000000
--- a/tests/nodes/argmax_fp16x16_1D_last_index/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(2);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp16x16_2D_keepdims_false.cairo b/tests/nodes/argmax_fp16x16_2D_keepdims_false.cairo
deleted file mode 100644
index 32fa5de64..000000000
--- a/tests/nodes/argmax_fp16x16_2D_keepdims_false.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::FP16x16TensorPartialEq;
-use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp16x16_2D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp16x16_2D_last_index.cairo b/tests/nodes/argmax_fp16x16_2D_last_index.cairo
deleted file mode 100644
index 3c0061e3e..000000000
--- a/tests/nodes/argmax_fp16x16_2D_last_index.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::FP16x16TensorPartialEq;
-use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp16x16_2D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp16x16_2D_last_index/output_0.cairo b/tests/nodes/argmax_fp16x16_2D_last_index/output_0.cairo
deleted file mode 100644
index a550af121..000000000
--- a/tests/nodes/argmax_fp16x16_2D_last_index/output_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp16x16_3D_default/input_0.cairo b/tests/nodes/argmax_fp16x16_3D_default/input_0.cairo
deleted file mode 100644
index d2f4ba6a0..000000000
--- a/tests/nodes/argmax_fp16x16_3D_default/input_0.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
-use orion::numbers::{FixedTrait, FP16x16};
-
-fn input_0() -> Tensor<FP16x16> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP16x16 { mag: 720896, sign: true });
-    data.append(FP16x16 { mag: 4390912, sign: false });
-    data.append(FP16x16 { mag: 5111808, sign: true });
-    data.append(FP16x16 { mag: 6553600, sign: false });
-    data.append(FP16x16 { mag: 1245184, sign: false });
-    data.append(FP16x16 { mag: 6094848, sign: true });
-    data.append(FP16x16 { mag: 4718592, sign: false });
-    data.append(FP16x16 { mag: 3211264, sign: true });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp16x16_3D_keepdims_false.cairo b/tests/nodes/argmax_fp16x16_3D_keepdims_false.cairo
deleted file mode 100644
index dfd529ffb..000000000
--- a/tests/nodes/argmax_fp16x16_3D_keepdims_false.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::FP16x16TensorPartialEq;
-use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp16x16_3D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp16x16_3D_keepdims_false/input_0.cairo b/tests/nodes/argmax_fp16x16_3D_keepdims_false/input_0.cairo
deleted file mode 100644
index 392ed4928..000000000
--- a/tests/nodes/argmax_fp16x16_3D_keepdims_false/input_0.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
-use orion::numbers::{FixedTrait, FP16x16};
-
-fn input_0() -> Tensor<FP16x16> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP16x16 { mag: 1572864, sign: false });
-    data.append(FP16x16 { mag: 2424832, sign: true });
-    data.append(FP16x16 { mag: 6422528, sign: false });
-    data.append(FP16x16 { mag: 5242880, sign: false });
-    data.append(FP16x16 { mag: 1703936, sign: true });
-    data.append(FP16x16 { mag: 6291456, sign: false });
-    data.append(FP16x16 { mag: 589824, sign: true });
-    data.append(FP16x16 { mag: 7733248, sign: true });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp16x16_3D_keepdims_false/output_0.cairo b/tests/nodes/argmax_fp16x16_3D_keepdims_false/output_0.cairo
deleted file mode 100644
index 825125d8a..000000000
--- a/tests/nodes/argmax_fp16x16_3D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,16 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(1);
-    data.append(0);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp16x16_3D_last_index.cairo b/tests/nodes/argmax_fp16x16_3D_last_index.cairo
deleted file mode 100644
index 72b0e2642..000000000
--- a/tests/nodes/argmax_fp16x16_3D_last_index.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::FP16x16TensorPartialEq;
-use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp16x16_3D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp16x16_3D_last_index/input_0.cairo b/tests/nodes/argmax_fp16x16_3D_last_index/input_0.cairo
deleted file mode 100644
index 50ac50fc8..000000000
--- a/tests/nodes/argmax_fp16x16_3D_last_index/input_0.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
-use orion::numbers::{FixedTrait, FP16x16};
-
-fn input_0() -> Tensor<FP16x16> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP16x16 { mag: 6291456, sign: true });
-    data.append(FP16x16 { mag: 7077888, sign: true });
-    data.append(FP16x16 { mag: 1966080, sign: false });
-    data.append(FP16x16 { mag: 4063232, sign: false });
-    data.append(FP16x16 { mag: 5439488, sign: false });
-    data.append(FP16x16 { mag: 1114112, sign: true });
-    data.append(FP16x16 { mag: 1441792, sign: false });
-    data.append(FP16x16 { mag: 458752, sign: true });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp16x16_3D_last_index/output_0.cairo b/tests/nodes/argmax_fp16x16_3D_last_index/output_0.cairo
deleted file mode 100644
index 9c34c6139..000000000
--- a/tests/nodes/argmax_fp16x16_3D_last_index/output_0.cairo
+++ /dev/null
@@ -1,17 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    data.append(1);
-    data.append(0);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_1D_default.cairo b/tests/nodes/argmax_fp8x23_1D_default.cairo
deleted file mode 100644
index 022e60b2f..000000000
--- a/tests/nodes/argmax_fp8x23_1D_default.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::operators::tensor::FP8x23TensorPartialEq;
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp8x23_1D_default() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp8x23_1D_default/input_0.cairo b/tests/nodes/argmax_fp8x23_1D_default/input_0.cairo
deleted file mode 100644
index 96ca7af0a..000000000
--- a/tests/nodes/argmax_fp8x23_1D_default/input_0.cairo
+++ /dev/null
@@ -1,15 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::numbers::{FixedTrait, FP8x23};
-
-fn input_0() -> Tensor<FP8x23> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP8x23 { mag: 1040187392, sign: false });
-    data.append(FP8x23 { mag: 637534208, sign: true });
-    data.append(FP8x23 { mag: 830472192, sign: false });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_1D_default/output_0.cairo b/tests/nodes/argmax_fp8x23_1D_default/output_0.cairo
deleted file mode 100644
index e4dbc507c..000000000
--- a/tests/nodes/argmax_fp8x23_1D_default/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_1D_keepdims_false.cairo b/tests/nodes/argmax_fp8x23_1D_keepdims_false.cairo
deleted file mode 100644
index cb6cea94f..000000000
--- a/tests/nodes/argmax_fp8x23_1D_keepdims_false.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::operators::tensor::FP8x23TensorPartialEq;
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp8x23_1D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp8x23_1D_keepdims_false/input_0.cairo b/tests/nodes/argmax_fp8x23_1D_keepdims_false/input_0.cairo
deleted file mode 100644
index 35f085f6e..000000000
--- a/tests/nodes/argmax_fp8x23_1D_keepdims_false/input_0.cairo
+++ /dev/null
@@ -1,15 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::numbers::{FixedTrait, FP8x23};
-
-fn input_0() -> Tensor<FP8x23> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP8x23 { mag: 402653184, sign: false });
-    data.append(FP8x23 { mag: 905969664, sign: true });
-    data.append(FP8x23 { mag: 822083584, sign: false });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_1D_keepdims_false/output_0.cairo b/tests/nodes/argmax_fp8x23_1D_keepdims_false/output_0.cairo
deleted file mode 100644
index 23846d97b..000000000
--- a/tests/nodes/argmax_fp8x23_1D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(2);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_1D_last_index.cairo b/tests/nodes/argmax_fp8x23_1D_last_index.cairo
deleted file mode 100644
index e2a7f0985..000000000
--- a/tests/nodes/argmax_fp8x23_1D_last_index.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::operators::tensor::FP8x23TensorPartialEq;
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp8x23_1D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp8x23_1D_last_index/input_0.cairo b/tests/nodes/argmax_fp8x23_1D_last_index/input_0.cairo
deleted file mode 100644
index 7613acbd4..000000000
--- a/tests/nodes/argmax_fp8x23_1D_last_index/input_0.cairo
+++ /dev/null
@@ -1,15 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::numbers::{FixedTrait, FP8x23};
-
-fn input_0() -> Tensor<FP8x23> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP8x23 { mag: 369098752, sign: false });
-    data.append(FP8x23 { mag: 847249408, sign: true });
-    data.append(FP8x23 { mag: 612368384, sign: false });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_1D_last_index/output_0.cairo b/tests/nodes/argmax_fp8x23_1D_last_index/output_0.cairo
deleted file mode 100644
index 23846d97b..000000000
--- a/tests/nodes/argmax_fp8x23_1D_last_index/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(2);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_2D_default.cairo b/tests/nodes/argmax_fp8x23_2D_default.cairo
deleted file mode 100644
index 0e474f982..000000000
--- a/tests/nodes/argmax_fp8x23_2D_default.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::operators::tensor::FP8x23TensorPartialEq;
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp8x23_2D_default() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp8x23_2D_default/input_0.cairo b/tests/nodes/argmax_fp8x23_2D_default/input_0.cairo
deleted file mode 100644
index d3ddea114..000000000
--- a/tests/nodes/argmax_fp8x23_2D_default/input_0.cairo
+++ /dev/null
@@ -1,17 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::numbers::{FixedTrait, FP8x23};
-
-fn input_0() -> Tensor<FP8x23> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP8x23 { mag: 654311424, sign: true });
-    data.append(FP8x23 { mag: 469762048, sign: false });
-    data.append(FP8x23 { mag: 964689920, sign: true });
-    data.append(FP8x23 { mag: 662700032, sign: false });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_2D_keepdims_false.cairo b/tests/nodes/argmax_fp8x23_2D_keepdims_false.cairo
deleted file mode 100644
index 40fe14c02..000000000
--- a/tests/nodes/argmax_fp8x23_2D_keepdims_false.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::operators::tensor::FP8x23TensorPartialEq;
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp8x23_2D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp8x23_2D_keepdims_false/input_0.cairo b/tests/nodes/argmax_fp8x23_2D_keepdims_false/input_0.cairo
deleted file mode 100644
index efb3441d3..000000000
--- a/tests/nodes/argmax_fp8x23_2D_keepdims_false/input_0.cairo
+++ /dev/null
@@ -1,17 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::numbers::{FixedTrait, FP8x23};
-
-fn input_0() -> Tensor<FP8x23> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP8x23 { mag: 536870912, sign: true });
-    data.append(FP8x23 { mag: 469762048, sign: false });
-    data.append(FP8x23 { mag: 83886080, sign: true });
-    data.append(FP8x23 { mag: 620756992, sign: true });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_2D_keepdims_false/output_0.cairo b/tests/nodes/argmax_fp8x23_2D_keepdims_false/output_0.cairo
deleted file mode 100644
index edc68637b..000000000
--- a/tests/nodes/argmax_fp8x23_2D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,13 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_2D_last_index.cairo b/tests/nodes/argmax_fp8x23_2D_last_index.cairo
deleted file mode 100644
index 05510e4f2..000000000
--- a/tests/nodes/argmax_fp8x23_2D_last_index.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::operators::tensor::FP8x23TensorPartialEq;
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp8x23_2D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp8x23_2D_last_index/input_0.cairo b/tests/nodes/argmax_fp8x23_2D_last_index/input_0.cairo
deleted file mode 100644
index be2e02d4c..000000000
--- a/tests/nodes/argmax_fp8x23_2D_last_index/input_0.cairo
+++ /dev/null
@@ -1,17 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::numbers::{FixedTrait, FP8x23};
-
-fn input_0() -> Tensor<FP8x23> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP8x23 { mag: 335544320, sign: false });
-    data.append(FP8x23 { mag: 1031798784, sign: true });
-    data.append(FP8x23 { mag: 989855744, sign: true });
-    data.append(FP8x23 { mag: 813694976, sign: true });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_2D_last_index/output_0.cairo b/tests/nodes/argmax_fp8x23_2D_last_index/output_0.cairo
deleted file mode 100644
index a550af121..000000000
--- a/tests/nodes/argmax_fp8x23_2D_last_index/output_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_3D_default.cairo b/tests/nodes/argmax_fp8x23_3D_default.cairo
deleted file mode 100644
index 4ca1b45c4..000000000
--- a/tests/nodes/argmax_fp8x23_3D_default.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::operators::tensor::FP8x23TensorPartialEq;
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp8x23_3D_default() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp8x23_3D_default/input_0.cairo b/tests/nodes/argmax_fp8x23_3D_default/input_0.cairo
deleted file mode 100644
index 4512a7d3c..000000000
--- a/tests/nodes/argmax_fp8x23_3D_default/input_0.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::numbers::{FixedTrait, FP8x23};
-
-fn input_0() -> Tensor<FP8x23> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP8x23 { mag: 612368384, sign: false });
-    data.append(FP8x23 { mag: 578813952, sign: false });
-    data.append(FP8x23 { mag: 947912704, sign: false });
-    data.append(FP8x23 { mag: 201326592, sign: true });
-    data.append(FP8x23 { mag: 1031798784, sign: true });
-    data.append(FP8x23 { mag: 729808896, sign: false });
-    data.append(FP8x23 { mag: 922746880, sign: false });
-    data.append(FP8x23 { mag: 33554432, sign: true });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_3D_default/output_0.cairo b/tests/nodes/argmax_fp8x23_3D_default/output_0.cairo
deleted file mode 100644
index 07cf1b47e..000000000
--- a/tests/nodes/argmax_fp8x23_3D_default/output_0.cairo
+++ /dev/null
@@ -1,17 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(1);
-    data.append(0);
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_3D_keepdims_false.cairo b/tests/nodes/argmax_fp8x23_3D_keepdims_false.cairo
deleted file mode 100644
index 4d5a5b850..000000000
--- a/tests/nodes/argmax_fp8x23_3D_keepdims_false.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::operators::tensor::FP8x23TensorPartialEq;
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp8x23_3D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp8x23_3D_keepdims_false/input_0.cairo b/tests/nodes/argmax_fp8x23_3D_keepdims_false/input_0.cairo
deleted file mode 100644
index a7ac816c9..000000000
--- a/tests/nodes/argmax_fp8x23_3D_keepdims_false/input_0.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::numbers::{FixedTrait, FP8x23};
-
-fn input_0() -> Tensor<FP8x23> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP8x23 { mag: 444596224, sign: false });
-    data.append(FP8x23 { mag: 369098752, sign: false });
-    data.append(FP8x23 { mag: 1056964608, sign: false });
-    data.append(FP8x23 { mag: 234881024, sign: false });
-    data.append(FP8x23 { mag: 159383552, sign: true });
-    data.append(FP8x23 { mag: 16777216, sign: true });
-    data.append(FP8x23 { mag: 897581056, sign: true });
-    data.append(FP8x23 { mag: 327155712, sign: true });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_3D_keepdims_false/output_0.cairo b/tests/nodes/argmax_fp8x23_3D_keepdims_false/output_0.cairo
deleted file mode 100644
index 637e10c01..000000000
--- a/tests/nodes/argmax_fp8x23_3D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,16 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(0);
-    data.append(0);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_3D_last_index.cairo b/tests/nodes/argmax_fp8x23_3D_last_index.cairo
deleted file mode 100644
index 856fbd988..000000000
--- a/tests/nodes/argmax_fp8x23_3D_last_index.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::operators::tensor::FP8x23TensorPartialEq;
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_fp8x23_3D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_fp8x23_3D_last_index/input_0.cairo b/tests/nodes/argmax_fp8x23_3D_last_index/input_0.cairo
deleted file mode 100644
index 42b2e25d0..000000000
--- a/tests/nodes/argmax_fp8x23_3D_last_index/input_0.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{FP8x23Tensor, FP8x23TensorAdd};
-use orion::numbers::{FixedTrait, FP8x23};
-
-fn input_0() -> Tensor<FP8x23> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(FP8x23 { mag: 83886080, sign: false });
-    data.append(FP8x23 { mag: 729808896, sign: true });
-    data.append(FP8x23 { mag: 536870912, sign: true });
-    data.append(FP8x23 { mag: 905969664, sign: false });
-    data.append(FP8x23 { mag: 75497472, sign: true });
-    data.append(FP8x23 { mag: 377487360, sign: true });
-    data.append(FP8x23 { mag: 92274688, sign: true });
-    data.append(FP8x23 { mag: 209715200, sign: true });
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp8x23_3D_last_index/output_0.cairo b/tests/nodes/argmax_fp8x23_3D_last_index/output_0.cairo
deleted file mode 100644
index e35cd05a2..000000000
--- a/tests/nodes/argmax_fp8x23_3D_last_index/output_0.cairo
+++ /dev/null
@@ -1,17 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(1);
-    data.append(1);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_1D_default.cairo b/tests/nodes/argmax_i32_1D_default.cairo
deleted file mode 100644
index 7846917bd..000000000
--- a/tests/nodes/argmax_i32_1D_default.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::I32TensorPartialEq;
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i32_1D_default() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i32_1D_default/output_0.cairo b/tests/nodes/argmax_i32_1D_default/output_0.cairo
deleted file mode 100644
index 23846d97b..000000000
--- a/tests/nodes/argmax_i32_1D_default/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(2);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_1D_keepdims_false.cairo b/tests/nodes/argmax_i32_1D_keepdims_false.cairo
deleted file mode 100644
index 9649435c7..000000000
--- a/tests/nodes/argmax_i32_1D_keepdims_false.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::I32TensorPartialEq;
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i32_1D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i32_1D_keepdims_false/input_0.cairo b/tests/nodes/argmax_i32_1D_keepdims_false/input_0.cairo
deleted file mode 100644
index 7a2414371..000000000
--- a/tests/nodes/argmax_i32_1D_keepdims_false/input_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-
-fn input_0() -> Tensor<i32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(-69);
-    data.append(51);
-    data.append(-49);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_1D_keepdims_false/output_0.cairo b/tests/nodes/argmax_i32_1D_keepdims_false/output_0.cairo
deleted file mode 100644
index f59edfe99..000000000
--- a/tests/nodes/argmax_i32_1D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_1D_last_index.cairo b/tests/nodes/argmax_i32_1D_last_index.cairo
deleted file mode 100644
index 94f533be6..000000000
--- a/tests/nodes/argmax_i32_1D_last_index.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::I32TensorPartialEq;
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i32_1D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i32_1D_last_index/input_0.cairo b/tests/nodes/argmax_i32_1D_last_index/input_0.cairo
deleted file mode 100644
index fa967cdc9..000000000
--- a/tests/nodes/argmax_i32_1D_last_index/input_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-
-fn input_0() -> Tensor<i32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(135);
-    data.append(148);
-    data.append(67);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_1D_last_index/output_0.cairo b/tests/nodes/argmax_i32_1D_last_index/output_0.cairo
deleted file mode 100644
index f59edfe99..000000000
--- a/tests/nodes/argmax_i32_1D_last_index/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_2D_default.cairo b/tests/nodes/argmax_i32_2D_default.cairo
deleted file mode 100644
index 1d8e4b3f7..000000000
--- a/tests/nodes/argmax_i32_2D_default.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::I32TensorPartialEq;
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i32_2D_default() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i32_2D_default/output_0.cairo b/tests/nodes/argmax_i32_2D_default/output_0.cairo
deleted file mode 100644
index e072284af..000000000
--- a/tests/nodes/argmax_i32_2D_default/output_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_2D_keepdims_false.cairo b/tests/nodes/argmax_i32_2D_keepdims_false.cairo
deleted file mode 100644
index 6e5372c2b..000000000
--- a/tests/nodes/argmax_i32_2D_keepdims_false.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::I32TensorPartialEq;
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i32_2D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i32_2D_keepdims_false/output_0.cairo b/tests/nodes/argmax_i32_2D_keepdims_false/output_0.cairo
deleted file mode 100644
index 4870d2de9..000000000
--- a/tests/nodes/argmax_i32_2D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,13 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_2D_last_index.cairo b/tests/nodes/argmax_i32_2D_last_index.cairo
deleted file mode 100644
index 5d734d434..000000000
--- a/tests/nodes/argmax_i32_2D_last_index.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::I32TensorPartialEq;
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i32_2D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i32_2D_last_index/output_0.cairo b/tests/nodes/argmax_i32_2D_last_index/output_0.cairo
deleted file mode 100644
index e072284af..000000000
--- a/tests/nodes/argmax_i32_2D_last_index/output_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_3D_default.cairo b/tests/nodes/argmax_i32_3D_default.cairo
deleted file mode 100644
index 0f1088869..000000000
--- a/tests/nodes/argmax_i32_3D_default.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::I32TensorPartialEq;
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i32_3D_default() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i32_3D_default/input_0.cairo b/tests/nodes/argmax_i32_3D_default/input_0.cairo
deleted file mode 100644
index 142afb3c3..000000000
--- a/tests/nodes/argmax_i32_3D_default/input_0.cairo
+++ /dev/null
@@ -1,21 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-
-fn input_0() -> Tensor<i32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(-56);
-    data.append(46);
-    data.append(-48);
-    data.append(91);
-    data.append(-70);
-    data.append(-54);
-    data.append(96);
-    data.append(122);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_3D_default/output_0.cairo b/tests/nodes/argmax_i32_3D_default/output_0.cairo
deleted file mode 100644
index a4cc4f838..000000000
--- a/tests/nodes/argmax_i32_3D_default/output_0.cairo
+++ /dev/null
@@ -1,17 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(0);
-    data.append(1);
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_3D_keepdims_false.cairo b/tests/nodes/argmax_i32_3D_keepdims_false.cairo
deleted file mode 100644
index 8aac46a2c..000000000
--- a/tests/nodes/argmax_i32_3D_keepdims_false.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::I32TensorPartialEq;
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i32_3D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i32_3D_keepdims_false/input_0.cairo b/tests/nodes/argmax_i32_3D_keepdims_false/input_0.cairo
deleted file mode 100644
index 5a61a206a..000000000
--- a/tests/nodes/argmax_i32_3D_keepdims_false/input_0.cairo
+++ /dev/null
@@ -1,21 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-
-fn input_0() -> Tensor<i32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(-121);
-    data.append(-114);
-    data.append(56);
-    data.append(42);
-    data.append(79);
-    data.append(43);
-    data.append(126);
-    data.append(-64);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_3D_keepdims_false/output_0.cairo b/tests/nodes/argmax_i32_3D_keepdims_false/output_0.cairo
deleted file mode 100644
index 1d637338a..000000000
--- a/tests/nodes/argmax_i32_3D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,16 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    data.append(1);
-    data.append(1);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_3D_last_index.cairo b/tests/nodes/argmax_i32_3D_last_index.cairo
deleted file mode 100644
index 81f4997de..000000000
--- a/tests/nodes/argmax_i32_3D_last_index.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::I32TensorPartialEq;
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i32_3D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i32_3D_last_index/input_0.cairo b/tests/nodes/argmax_i32_3D_last_index/input_0.cairo
deleted file mode 100644
index c5d8bc74e..000000000
--- a/tests/nodes/argmax_i32_3D_last_index/input_0.cairo
+++ /dev/null
@@ -1,21 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I32Tensor, I32TensorAdd};
-
-fn input_0() -> Tensor<i32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(98);
-    data.append(89);
-    data.append(-126);
-    data.append(-68);
-    data.append(31);
-    data.append(7);
-    data.append(-86);
-    data.append(99);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i32_3D_last_index/output_0.cairo b/tests/nodes/argmax_i32_3D_last_index/output_0.cairo
deleted file mode 100644
index a4cc4f838..000000000
--- a/tests/nodes/argmax_i32_3D_last_index/output_0.cairo
+++ /dev/null
@@ -1,17 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(0);
-    data.append(1);
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_1D_default.cairo b/tests/nodes/argmax_i8_1D_default.cairo
deleted file mode 100644
index 5fd40f6b5..000000000
--- a/tests/nodes/argmax_i8_1D_default.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::I8TensorPartialEq;
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i8_1D_default() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i8_1D_default/input_0.cairo b/tests/nodes/argmax_i8_1D_default/input_0.cairo
deleted file mode 100644
index d760a1faa..000000000
--- a/tests/nodes/argmax_i8_1D_default/input_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-
-fn input_0() -> Tensor<i8> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(6);
-    data.append(-60);
-    data.append(-123);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_1D_default/output_0.cairo b/tests/nodes/argmax_i8_1D_default/output_0.cairo
deleted file mode 100644
index e4dbc507c..000000000
--- a/tests/nodes/argmax_i8_1D_default/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_1D_keepdims_false.cairo b/tests/nodes/argmax_i8_1D_keepdims_false.cairo
deleted file mode 100644
index 9fa8ef88c..000000000
--- a/tests/nodes/argmax_i8_1D_keepdims_false.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::I8TensorPartialEq;
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i8_1D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i8_1D_keepdims_false/input_0.cairo b/tests/nodes/argmax_i8_1D_keepdims_false/input_0.cairo
deleted file mode 100644
index fdfdf6ff4..000000000
--- a/tests/nodes/argmax_i8_1D_keepdims_false/input_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-
-fn input_0() -> Tensor<i8> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(-68);
-    data.append(56);
-    data.append(-126);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_1D_keepdims_false/output_0.cairo b/tests/nodes/argmax_i8_1D_keepdims_false/output_0.cairo
deleted file mode 100644
index f59edfe99..000000000
--- a/tests/nodes/argmax_i8_1D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_1D_last_index.cairo b/tests/nodes/argmax_i8_1D_last_index.cairo
deleted file mode 100644
index 54ccb82f9..000000000
--- a/tests/nodes/argmax_i8_1D_last_index.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::I8TensorPartialEq;
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i8_1D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i8_1D_last_index/input_0.cairo b/tests/nodes/argmax_i8_1D_last_index/input_0.cairo
deleted file mode 100644
index 262a8306c..000000000
--- a/tests/nodes/argmax_i8_1D_last_index/input_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-
-fn input_0() -> Tensor<i8> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(111);
-    data.append(84);
-    data.append(-120);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_1D_last_index/output_0.cairo b/tests/nodes/argmax_i8_1D_last_index/output_0.cairo
deleted file mode 100644
index e4dbc507c..000000000
--- a/tests/nodes/argmax_i8_1D_last_index/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_2D_default.cairo b/tests/nodes/argmax_i8_2D_default.cairo
deleted file mode 100644
index dc0a71e4e..000000000
--- a/tests/nodes/argmax_i8_2D_default.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::I8TensorPartialEq;
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i8_2D_default() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i8_2D_default/input_0.cairo b/tests/nodes/argmax_i8_2D_default/input_0.cairo
deleted file mode 100644
index 9aa9650c7..000000000
--- a/tests/nodes/argmax_i8_2D_default/input_0.cairo
+++ /dev/null
@@ -1,16 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-
-fn input_0() -> Tensor<i8> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(-8);
-    data.append(-67);
-    data.append(-22);
-    data.append(99);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_2D_default/output_0.cairo b/tests/nodes/argmax_i8_2D_default/output_0.cairo
deleted file mode 100644
index a550af121..000000000
--- a/tests/nodes/argmax_i8_2D_default/output_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_2D_keepdims_false.cairo b/tests/nodes/argmax_i8_2D_keepdims_false.cairo
deleted file mode 100644
index a56d67f01..000000000
--- a/tests/nodes/argmax_i8_2D_keepdims_false.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::I8TensorPartialEq;
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i8_2D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i8_2D_keepdims_false/input_0.cairo b/tests/nodes/argmax_i8_2D_keepdims_false/input_0.cairo
deleted file mode 100644
index 373558d6e..000000000
--- a/tests/nodes/argmax_i8_2D_keepdims_false/input_0.cairo
+++ /dev/null
@@ -1,16 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-
-fn input_0() -> Tensor<i8> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(25);
-    data.append(-120);
-    data.append(-2);
-    data.append(105);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_2D_keepdims_false/output_0.cairo b/tests/nodes/argmax_i8_2D_keepdims_false/output_0.cairo
deleted file mode 100644
index 2fe75c127..000000000
--- a/tests/nodes/argmax_i8_2D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,13 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_2D_last_index.cairo b/tests/nodes/argmax_i8_2D_last_index.cairo
deleted file mode 100644
index 1574a72f7..000000000
--- a/tests/nodes/argmax_i8_2D_last_index.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::I8TensorPartialEq;
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i8_2D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i8_2D_last_index/input_0.cairo b/tests/nodes/argmax_i8_2D_last_index/input_0.cairo
deleted file mode 100644
index 051786fb8..000000000
--- a/tests/nodes/argmax_i8_2D_last_index/input_0.cairo
+++ /dev/null
@@ -1,16 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-
-fn input_0() -> Tensor<i8> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(-73);
-    data.append(111);
-    data.append(47);
-    data.append(-18);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_2D_last_index/output_0.cairo b/tests/nodes/argmax_i8_2D_last_index/output_0.cairo
deleted file mode 100644
index 221e8a246..000000000
--- a/tests/nodes/argmax_i8_2D_last_index/output_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_3D_default.cairo b/tests/nodes/argmax_i8_3D_default.cairo
deleted file mode 100644
index eaf92eae5..000000000
--- a/tests/nodes/argmax_i8_3D_default.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::I8TensorPartialEq;
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i8_3D_default() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i8_3D_default/input_0.cairo b/tests/nodes/argmax_i8_3D_default/input_0.cairo
deleted file mode 100644
index cc96b52ca..000000000
--- a/tests/nodes/argmax_i8_3D_default/input_0.cairo
+++ /dev/null
@@ -1,21 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-
-fn input_0() -> Tensor<i8> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(66);
-    data.append(-56);
-    data.append(49);
-    data.append(-2);
-    data.append(-93);
-    data.append(-55);
-    data.append(115);
-    data.append(28);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_3D_default/output_0.cairo b/tests/nodes/argmax_i8_3D_default/output_0.cairo
deleted file mode 100644
index 56d0d331d..000000000
--- a/tests/nodes/argmax_i8_3D_default/output_0.cairo
+++ /dev/null
@@ -1,17 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(1);
-    data.append(1);
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_3D_keepdims_false.cairo b/tests/nodes/argmax_i8_3D_keepdims_false.cairo
deleted file mode 100644
index 3b3f3c10c..000000000
--- a/tests/nodes/argmax_i8_3D_keepdims_false.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::I8TensorPartialEq;
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i8_3D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i8_3D_keepdims_false/input_0.cairo b/tests/nodes/argmax_i8_3D_keepdims_false/input_0.cairo
deleted file mode 100644
index f6de84eca..000000000
--- a/tests/nodes/argmax_i8_3D_keepdims_false/input_0.cairo
+++ /dev/null
@@ -1,21 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-
-fn input_0() -> Tensor<i8> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(82);
-    data.append(70);
-    data.append(-18);
-    data.append(-46);
-    data.append(-4);
-    data.append(82);
-    data.append(58);
-    data.append(-41);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_3D_keepdims_false/output_0.cairo b/tests/nodes/argmax_i8_3D_keepdims_false/output_0.cairo
deleted file mode 100644
index f13678224..000000000
--- a/tests/nodes/argmax_i8_3D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,16 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(1);
-    data.append(1);
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_3D_last_index.cairo b/tests/nodes/argmax_i8_3D_last_index.cairo
deleted file mode 100644
index cd460d7c2..000000000
--- a/tests/nodes/argmax_i8_3D_last_index.cairo
+++ /dev/null
@@ -1,22 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::I8TensorPartialEq;
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_i8_3D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_i8_3D_last_index/input_0.cairo b/tests/nodes/argmax_i8_3D_last_index/input_0.cairo
deleted file mode 100644
index a665d1534..000000000
--- a/tests/nodes/argmax_i8_3D_last_index/input_0.cairo
+++ /dev/null
@@ -1,21 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{I8Tensor, I8TensorAdd};
-
-fn input_0() -> Tensor<i8> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(-7);
-    data.append(-94);
-    data.append(-64);
-    data.append(-19);
-    data.append(59);
-    data.append(-40);
-    data.append(99);
-    data.append(38);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_i8_3D_last_index/output_0.cairo b/tests/nodes/argmax_i8_3D_last_index/output_0.cairo
deleted file mode 100644
index 23b7ccc3c..000000000
--- a/tests/nodes/argmax_i8_3D_last_index/output_0.cairo
+++ /dev/null
@@ -1,17 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    data.append(1);
-    data.append(1);
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_fp16x16_1D_default.cairo b/tests/nodes/argmax_keepdims.cairo
similarity index 58%
rename from tests/nodes/argmax_fp16x16_1D_default.cairo
rename to tests/nodes/argmax_keepdims.cairo
index 873579cce..49f6f5d7b 100644
--- a/tests/nodes/argmax_fp16x16_1D_default.cairo
+++ b/tests/nodes/argmax_keepdims.cairo
@@ -2,21 +2,21 @@ mod input_0;
 mod output_0;
 
 
+use orion::operators::tensor::{I32Tensor, I32TensorAdd};
 use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
 use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::FP16x16TensorPartialEq;
+use orion::utils::{assert_eq, assert_seq_eq};
+use orion::operators::tensor::I32TensorPartialEq;
 use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
+use orion::operators::tensor::FP16x16TensorPartialEq;
 
 #[test]
 #[available_gas(2000000000)]
-fn test_argmax_fp16x16_1D_default() {
+fn test_argmax_keepdims() {
     let input_0 = input_0::input_0();
-    let z = output_0::output_0();
+    let z_0 = output_0::output_0();
 
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
+    let y_0 = input_0.argmax(1, Option::Some(true), Option::None(()));
 
-    assert_eq(y, z);
+    assert_eq(y_0, z_0);
 }
diff --git a/tests/nodes/argmax_fp16x16_2D_default/input_0.cairo b/tests/nodes/argmax_keepdims/input_0.cairo
similarity index 65%
rename from tests/nodes/argmax_fp16x16_2D_default/input_0.cairo
rename to tests/nodes/argmax_keepdims/input_0.cairo
index ce1b11cae..2bd0e9ded 100644
--- a/tests/nodes/argmax_fp16x16_2D_default/input_0.cairo
+++ b/tests/nodes/argmax_keepdims/input_0.cairo
@@ -9,9 +9,9 @@ fn input_0() -> Tensor<FP16x16> {
     shape.append(2);
 
     let mut data = ArrayTrait::new();
-    data.append(FP16x16 { mag: 196608, sign: true });
-    data.append(FP16x16 { mag: 2162688, sign: false });
-    data.append(FP16x16 { mag: 1507328, sign: true });
-    data.append(FP16x16 { mag: 6815744, sign: false });
+    data.append(FP16x16 { mag: 2, sign: false });
+    data.append(FP16x16 { mag: 1, sign: false });
+    data.append(FP16x16 { mag: 3, sign: false });
+    data.append(FP16x16 { mag: 10, sign: false });
     TensorTrait::new(shape.span(), data.span())
 }
diff --git a/tests/nodes/argmax_fp8x23_2D_default/output_0.cairo b/tests/nodes/argmax_keepdims/output_0.cairo
similarity index 72%
rename from tests/nodes/argmax_fp8x23_2D_default/output_0.cairo
rename to tests/nodes/argmax_keepdims/output_0.cairo
index a550af121..a1ae6cd36 100644
--- a/tests/nodes/argmax_fp8x23_2D_default/output_0.cairo
+++ b/tests/nodes/argmax_keepdims/output_0.cairo
@@ -1,11 +1,12 @@
 use core::array::{ArrayTrait, SpanTrait};
 use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
+use orion::operators::tensor::{I32Tensor, I32TensorAdd};
+use orion::numbers::NumberTrait;
 
-fn output_0() -> Tensor<u32> {
+fn output_0() -> Tensor<i32> {
     let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
     shape.append(2);
+    shape.append(1);
 
     let mut data = ArrayTrait::new();
     data.append(0);
diff --git a/tests/nodes/argmax_fp16x16_1D_keepdims_false.cairo b/tests/nodes/argmax_keepdims_select_last_index.cairo
similarity index 56%
rename from tests/nodes/argmax_fp16x16_1D_keepdims_false.cairo
rename to tests/nodes/argmax_keepdims_select_last_index.cairo
index 1563fad06..234cb0057 100644
--- a/tests/nodes/argmax_fp16x16_1D_keepdims_false.cairo
+++ b/tests/nodes/argmax_keepdims_select_last_index.cairo
@@ -2,21 +2,21 @@ mod input_0;
 mod output_0;
 
 
+use orion::operators::tensor::{I32Tensor, I32TensorAdd};
 use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
 use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::FP16x16TensorPartialEq;
+use orion::utils::{assert_eq, assert_seq_eq};
+use orion::operators::tensor::I32TensorPartialEq;
 use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
+use orion::operators::tensor::FP16x16TensorPartialEq;
 
 #[test]
 #[available_gas(2000000000)]
-fn test_argmax_fp16x16_1D_keepdims_false() {
+fn test_argmax_keepdims_select_last_index() {
     let input_0 = input_0::input_0();
-    let z = output_0::output_0();
+    let z_0 = output_0::output_0();
 
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
+    let y_0 = input_0.argmax(1, Option::Some(true), Option::Some(true));
 
-    assert_eq(y, z);
+    assert_eq(y_0, z_0);
 }
diff --git a/tests/nodes/argmax_fp16x16_1D_keepdims_false/input_0.cairo b/tests/nodes/argmax_keepdims_select_last_index/input_0.cairo
similarity index 60%
rename from tests/nodes/argmax_fp16x16_1D_keepdims_false/input_0.cairo
rename to tests/nodes/argmax_keepdims_select_last_index/input_0.cairo
index 1c4977212..c0673e50d 100644
--- a/tests/nodes/argmax_fp16x16_1D_keepdims_false/input_0.cairo
+++ b/tests/nodes/argmax_keepdims_select_last_index/input_0.cairo
@@ -5,11 +5,13 @@ use orion::numbers::{FixedTrait, FP16x16};
 
 fn input_0() -> Tensor<FP16x16> {
     let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
+    shape.append(2);
+    shape.append(2);
 
     let mut data = ArrayTrait::new();
-    data.append(FP16x16 { mag: 6422528, sign: false });
-    data.append(FP16x16 { mag: 2031616, sign: false });
-    data.append(FP16x16 { mag: 6356992, sign: true });
+    data.append(FP16x16 { mag: 2, sign: false });
+    data.append(FP16x16 { mag: 2, sign: false });
+    data.append(FP16x16 { mag: 3, sign: false });
+    data.append(FP16x16 { mag: 10, sign: false });
     TensorTrait::new(shape.span(), data.span())
 }
diff --git a/tests/nodes/argmax_i32_2D_default/input_0.cairo b/tests/nodes/argmax_keepdims_select_last_index/output_0.cairo
similarity index 68%
rename from tests/nodes/argmax_i32_2D_default/input_0.cairo
rename to tests/nodes/argmax_keepdims_select_last_index/output_0.cairo
index 92434e2c9..241b14b08 100644
--- a/tests/nodes/argmax_i32_2D_default/input_0.cairo
+++ b/tests/nodes/argmax_keepdims_select_last_index/output_0.cairo
@@ -1,16 +1,15 @@
 use core::array::{ArrayTrait, SpanTrait};
 use orion::operators::tensor::{TensorTrait, Tensor};
 use orion::operators::tensor::{I32Tensor, I32TensorAdd};
+use orion::numbers::NumberTrait;
 
-fn input_0() -> Tensor<i32> {
+fn output_0() -> Tensor<i32> {
     let mut shape = ArrayTrait::<usize>::new();
     shape.append(2);
-    shape.append(2);
+    shape.append(1);
 
     let mut data = ArrayTrait::new();
-    data.append(-27);
-    data.append(114);
-    data.append(-54);
-    data.append(-100);
+    data.append(1);
+    data.append(1);
     TensorTrait::new(shape.span(), data.span())
 }
diff --git a/tests/nodes/argmax_fp16x16_2D_default.cairo b/tests/nodes/argmax_negative_axis_keepdims.cairo
similarity index 56%
rename from tests/nodes/argmax_fp16x16_2D_default.cairo
rename to tests/nodes/argmax_negative_axis_keepdims.cairo
index 754dd2119..85eb2619f 100644
--- a/tests/nodes/argmax_fp16x16_2D_default.cairo
+++ b/tests/nodes/argmax_negative_axis_keepdims.cairo
@@ -2,21 +2,21 @@ mod input_0;
 mod output_0;
 
 
+use orion::operators::tensor::{I32Tensor, I32TensorAdd};
 use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
 use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::FP16x16TensorPartialEq;
+use orion::utils::{assert_eq, assert_seq_eq};
+use orion::operators::tensor::I32TensorPartialEq;
 use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
+use orion::operators::tensor::FP16x16TensorPartialEq;
 
 #[test]
 #[available_gas(2000000000)]
-fn test_argmax_fp16x16_2D_default() {
+fn test_argmax_negative_axis_keepdims() {
     let input_0 = input_0::input_0();
-    let z = output_0::output_0();
+    let z_0 = output_0::output_0();
 
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
+    let y_0 = input_0.argmax(-1, Option::Some(true), Option::None(()));
 
-    assert_eq(y, z);
+    assert_eq(y_0, z_0);
 }
diff --git a/tests/nodes/argmax_negative_axis_keepdims/input_0.cairo b/tests/nodes/argmax_negative_axis_keepdims/input_0.cairo
new file mode 100644
index 000000000..2bd0e9ded
--- /dev/null
+++ b/tests/nodes/argmax_negative_axis_keepdims/input_0.cairo
@@ -0,0 +1,17 @@
+use core::array::{ArrayTrait, SpanTrait};
+use orion::operators::tensor::{TensorTrait, Tensor};
+use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
+use orion::numbers::{FixedTrait, FP16x16};
+
+fn input_0() -> Tensor<FP16x16> {
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(2);
+    shape.append(2);
+
+    let mut data = ArrayTrait::new();
+    data.append(FP16x16 { mag: 2, sign: false });
+    data.append(FP16x16 { mag: 1, sign: false });
+    data.append(FP16x16 { mag: 3, sign: false });
+    data.append(FP16x16 { mag: 10, sign: false });
+    TensorTrait::new(shape.span(), data.span())
+}
diff --git a/tests/nodes/argmax_fp16x16_2D_default/output_0.cairo b/tests/nodes/argmax_negative_axis_keepdims/output_0.cairo
similarity index 72%
rename from tests/nodes/argmax_fp16x16_2D_default/output_0.cairo
rename to tests/nodes/argmax_negative_axis_keepdims/output_0.cairo
index a550af121..a1ae6cd36 100644
--- a/tests/nodes/argmax_fp16x16_2D_default/output_0.cairo
+++ b/tests/nodes/argmax_negative_axis_keepdims/output_0.cairo
@@ -1,11 +1,12 @@
 use core::array::{ArrayTrait, SpanTrait};
 use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
+use orion::operators::tensor::{I32Tensor, I32TensorAdd};
+use orion::numbers::NumberTrait;
 
-fn output_0() -> Tensor<u32> {
+fn output_0() -> Tensor<i32> {
     let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
     shape.append(2);
+    shape.append(1);
 
     let mut data = ArrayTrait::new();
     data.append(0);
diff --git a/tests/nodes/argmax_negative_axis_keepdims_select_last_index.cairo b/tests/nodes/argmax_negative_axis_keepdims_select_last_index.cairo
new file mode 100644
index 000000000..92c5f9f10
--- /dev/null
+++ b/tests/nodes/argmax_negative_axis_keepdims_select_last_index.cairo
@@ -0,0 +1,22 @@
+mod input_0;
+mod output_0;
+
+
+use orion::operators::tensor::{I32Tensor, I32TensorAdd};
+use core::array::{ArrayTrait, SpanTrait};
+use orion::operators::tensor::{TensorTrait, Tensor};
+use orion::utils::{assert_eq, assert_seq_eq};
+use orion::operators::tensor::I32TensorPartialEq;
+use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
+use orion::operators::tensor::FP16x16TensorPartialEq;
+
+#[test]
+#[available_gas(2000000000)]
+fn test_argmax_negative_axis_keepdims_select_last_index() {
+    let input_0 = input_0::input_0();
+    let z_0 = output_0::output_0();
+
+    let y_0 = input_0.argmax(-1, Option::Some(true), Option::Some(true));
+
+    assert_eq(y_0, z_0);
+}
diff --git a/tests/nodes/argmax_negative_axis_keepdims_select_last_index/input_0.cairo b/tests/nodes/argmax_negative_axis_keepdims_select_last_index/input_0.cairo
new file mode 100644
index 000000000..c0673e50d
--- /dev/null
+++ b/tests/nodes/argmax_negative_axis_keepdims_select_last_index/input_0.cairo
@@ -0,0 +1,17 @@
+use core::array::{ArrayTrait, SpanTrait};
+use orion::operators::tensor::{TensorTrait, Tensor};
+use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
+use orion::numbers::{FixedTrait, FP16x16};
+
+fn input_0() -> Tensor<FP16x16> {
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(2);
+    shape.append(2);
+
+    let mut data = ArrayTrait::new();
+    data.append(FP16x16 { mag: 2, sign: false });
+    data.append(FP16x16 { mag: 2, sign: false });
+    data.append(FP16x16 { mag: 3, sign: false });
+    data.append(FP16x16 { mag: 10, sign: false });
+    TensorTrait::new(shape.span(), data.span())
+}
diff --git a/tests/nodes/argmax_fp16x16_3D_default/output_0.cairo b/tests/nodes/argmax_negative_axis_keepdims_select_last_index/output_0.cairo
similarity index 67%
rename from tests/nodes/argmax_fp16x16_3D_default/output_0.cairo
rename to tests/nodes/argmax_negative_axis_keepdims_select_last_index/output_0.cairo
index 03945682b..241b14b08 100644
--- a/tests/nodes/argmax_fp16x16_3D_default/output_0.cairo
+++ b/tests/nodes/argmax_negative_axis_keepdims_select_last_index/output_0.cairo
@@ -1,17 +1,15 @@
 use core::array::{ArrayTrait, SpanTrait};
 use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
+use orion::operators::tensor::{I32Tensor, I32TensorAdd};
+use orion::numbers::NumberTrait;
 
-fn output_0() -> Tensor<u32> {
+fn output_0() -> Tensor<i32> {
     let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
     shape.append(2);
+    shape.append(1);
 
     let mut data = ArrayTrait::new();
     data.append(1);
-    data.append(0);
     data.append(1);
-    data.append(0);
     TensorTrait::new(shape.span(), data.span())
 }
diff --git a/tests/nodes/argmax_fp16x16_3D_default.cairo b/tests/nodes/argmax_no_keepdims.cairo
similarity index 57%
rename from tests/nodes/argmax_fp16x16_3D_default.cairo
rename to tests/nodes/argmax_no_keepdims.cairo
index b7efa083b..c278fee35 100644
--- a/tests/nodes/argmax_fp16x16_3D_default.cairo
+++ b/tests/nodes/argmax_no_keepdims.cairo
@@ -2,21 +2,21 @@ mod input_0;
 mod output_0;
 
 
+use orion::operators::tensor::{I32Tensor, I32TensorAdd};
 use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::utils::{assert_eq, assert_seq_eq};
 use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::FP16x16TensorPartialEq;
+use orion::utils::{assert_eq, assert_seq_eq};
+use orion::operators::tensor::I32TensorPartialEq;
 use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
+use orion::operators::tensor::FP16x16TensorPartialEq;
 
 #[test]
 #[available_gas(2000000000)]
-fn test_argmax_fp16x16_3D_default() {
+fn test_argmax_no_keepdims() {
     let input_0 = input_0::input_0();
-    let z = output_0::output_0();
+    let z_0 = output_0::output_0();
 
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
+    let y_0 = input_0.argmax(1, Option::Some(false), Option::None(()));
 
-    assert_eq(y, z);
+    assert_eq(y_0, z_0);
 }
diff --git a/tests/nodes/argmax_no_keepdims/input_0.cairo b/tests/nodes/argmax_no_keepdims/input_0.cairo
new file mode 100644
index 000000000..2bd0e9ded
--- /dev/null
+++ b/tests/nodes/argmax_no_keepdims/input_0.cairo
@@ -0,0 +1,17 @@
+use core::array::{ArrayTrait, SpanTrait};
+use orion::operators::tensor::{TensorTrait, Tensor};
+use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
+use orion::numbers::{FixedTrait, FP16x16};
+
+fn input_0() -> Tensor<FP16x16> {
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(2);
+    shape.append(2);
+
+    let mut data = ArrayTrait::new();
+    data.append(FP16x16 { mag: 2, sign: false });
+    data.append(FP16x16 { mag: 1, sign: false });
+    data.append(FP16x16 { mag: 3, sign: false });
+    data.append(FP16x16 { mag: 10, sign: false });
+    TensorTrait::new(shape.span(), data.span())
+}
diff --git a/tests/nodes/argmax_fp16x16_2D_keepdims_false/output_0.cairo b/tests/nodes/argmax_no_keepdims/output_0.cairo
similarity index 70%
rename from tests/nodes/argmax_fp16x16_2D_keepdims_false/output_0.cairo
rename to tests/nodes/argmax_no_keepdims/output_0.cairo
index 2fe75c127..ed50e53b0 100644
--- a/tests/nodes/argmax_fp16x16_2D_keepdims_false/output_0.cairo
+++ b/tests/nodes/argmax_no_keepdims/output_0.cairo
@@ -1,8 +1,9 @@
 use core::array::{ArrayTrait, SpanTrait};
 use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
+use orion::operators::tensor::{I32Tensor, I32TensorAdd};
+use orion::numbers::NumberTrait;
 
-fn output_0() -> Tensor<u32> {
+fn output_0() -> Tensor<i32> {
     let mut shape = ArrayTrait::<usize>::new();
     shape.append(2);
 
diff --git a/tests/nodes/argmax_no_keepdims_select_last_index.cairo b/tests/nodes/argmax_no_keepdims_select_last_index.cairo
new file mode 100644
index 000000000..5233204eb
--- /dev/null
+++ b/tests/nodes/argmax_no_keepdims_select_last_index.cairo
@@ -0,0 +1,22 @@
+mod input_0;
+mod output_0;
+
+
+use orion::operators::tensor::{I32Tensor, I32TensorAdd};
+use core::array::{ArrayTrait, SpanTrait};
+use orion::operators::tensor::{TensorTrait, Tensor};
+use orion::utils::{assert_eq, assert_seq_eq};
+use orion::operators::tensor::I32TensorPartialEq;
+use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
+use orion::operators::tensor::FP16x16TensorPartialEq;
+
+#[test]
+#[available_gas(2000000000)]
+fn test_argmax_no_keepdims_select_last_index() {
+    let input_0 = input_0::input_0();
+    let z_0 = output_0::output_0();
+
+    let y_0 = input_0.argmax(1, Option::Some(false), Option::Some(true));
+
+    assert_eq(y_0, z_0);
+}
diff --git a/tests/nodes/argmax_no_keepdims_select_last_index/input_0.cairo b/tests/nodes/argmax_no_keepdims_select_last_index/input_0.cairo
new file mode 100644
index 000000000..c0673e50d
--- /dev/null
+++ b/tests/nodes/argmax_no_keepdims_select_last_index/input_0.cairo
@@ -0,0 +1,17 @@
+use core::array::{ArrayTrait, SpanTrait};
+use orion::operators::tensor::{TensorTrait, Tensor};
+use orion::operators::tensor::{FP16x16Tensor, FP16x16TensorAdd};
+use orion::numbers::{FixedTrait, FP16x16};
+
+fn input_0() -> Tensor<FP16x16> {
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(2);
+    shape.append(2);
+
+    let mut data = ArrayTrait::new();
+    data.append(FP16x16 { mag: 2, sign: false });
+    data.append(FP16x16 { mag: 2, sign: false });
+    data.append(FP16x16 { mag: 3, sign: false });
+    data.append(FP16x16 { mag: 10, sign: false });
+    TensorTrait::new(shape.span(), data.span())
+}
diff --git a/tests/nodes/argmax_i32_1D_default/input_0.cairo b/tests/nodes/argmax_no_keepdims_select_last_index/output_0.cairo
similarity index 69%
rename from tests/nodes/argmax_i32_1D_default/input_0.cairo
rename to tests/nodes/argmax_no_keepdims_select_last_index/output_0.cairo
index 97fbd451a..29a921caf 100644
--- a/tests/nodes/argmax_i32_1D_default/input_0.cairo
+++ b/tests/nodes/argmax_no_keepdims_select_last_index/output_0.cairo
@@ -1,14 +1,14 @@
 use core::array::{ArrayTrait, SpanTrait};
 use orion::operators::tensor::{TensorTrait, Tensor};
 use orion::operators::tensor::{I32Tensor, I32TensorAdd};
+use orion::numbers::NumberTrait;
 
-fn input_0() -> Tensor<i32> {
+fn output_0() -> Tensor<i32> {
     let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
+    shape.append(2);
 
     let mut data = ArrayTrait::new();
-    data.append(-54);
-    data.append(29);
-    data.append(81);
+    data.append(1);
+    data.append(1);
     TensorTrait::new(shape.span(), data.span())
 }
diff --git a/tests/nodes/argmax_u32_1D_default.cairo b/tests/nodes/argmax_u32_1D_default.cairo
deleted file mode 100644
index f8bf64a5b..000000000
--- a/tests/nodes/argmax_u32_1D_default.cairo
+++ /dev/null
@@ -1,20 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_u32_1D_default() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_u32_1D_default/input_0.cairo b/tests/nodes/argmax_u32_1D_default/input_0.cairo
deleted file mode 100644
index d02f0a9a5..000000000
--- a/tests/nodes/argmax_u32_1D_default/input_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn input_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(244);
-    data.append(135);
-    data.append(53);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_1D_default/output_0.cairo b/tests/nodes/argmax_u32_1D_default/output_0.cairo
deleted file mode 100644
index e4dbc507c..000000000
--- a/tests/nodes/argmax_u32_1D_default/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_1D_keepdims_false.cairo b/tests/nodes/argmax_u32_1D_keepdims_false.cairo
deleted file mode 100644
index aa6555d1c..000000000
--- a/tests/nodes/argmax_u32_1D_keepdims_false.cairo
+++ /dev/null
@@ -1,20 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_u32_1D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_u32_1D_keepdims_false/input_0.cairo b/tests/nodes/argmax_u32_1D_keepdims_false/input_0.cairo
deleted file mode 100644
index 161945ea0..000000000
--- a/tests/nodes/argmax_u32_1D_keepdims_false/input_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn input_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(89);
-    data.append(33);
-    data.append(140);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_1D_keepdims_false/output_0.cairo b/tests/nodes/argmax_u32_1D_keepdims_false/output_0.cairo
deleted file mode 100644
index 23846d97b..000000000
--- a/tests/nodes/argmax_u32_1D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(2);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_1D_last_index.cairo b/tests/nodes/argmax_u32_1D_last_index.cairo
deleted file mode 100644
index e6ad1f332..000000000
--- a/tests/nodes/argmax_u32_1D_last_index.cairo
+++ /dev/null
@@ -1,20 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_u32_1D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_u32_1D_last_index/input_0.cairo b/tests/nodes/argmax_u32_1D_last_index/input_0.cairo
deleted file mode 100644
index caad0d75f..000000000
--- a/tests/nodes/argmax_u32_1D_last_index/input_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn input_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(3);
-
-    let mut data = ArrayTrait::new();
-    data.append(63);
-    data.append(40);
-    data.append(22);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_1D_last_index/output_0.cairo b/tests/nodes/argmax_u32_1D_last_index/output_0.cairo
deleted file mode 100644
index e4dbc507c..000000000
--- a/tests/nodes/argmax_u32_1D_last_index/output_0.cairo
+++ /dev/null
@@ -1,12 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_2D_default.cairo b/tests/nodes/argmax_u32_2D_default.cairo
deleted file mode 100644
index 9c9a285bd..000000000
--- a/tests/nodes/argmax_u32_2D_default.cairo
+++ /dev/null
@@ -1,20 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_u32_2D_default() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_u32_2D_default/input_0.cairo b/tests/nodes/argmax_u32_2D_default/input_0.cairo
deleted file mode 100644
index e6678379b..000000000
--- a/tests/nodes/argmax_u32_2D_default/input_0.cairo
+++ /dev/null
@@ -1,16 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn input_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(51);
-    data.append(202);
-    data.append(164);
-    data.append(83);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_2D_default/output_0.cairo b/tests/nodes/argmax_u32_2D_default/output_0.cairo
deleted file mode 100644
index 221e8a246..000000000
--- a/tests/nodes/argmax_u32_2D_default/output_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_2D_keepdims_false.cairo b/tests/nodes/argmax_u32_2D_keepdims_false.cairo
deleted file mode 100644
index 8bfa99454..000000000
--- a/tests/nodes/argmax_u32_2D_keepdims_false.cairo
+++ /dev/null
@@ -1,20 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_u32_2D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_u32_2D_keepdims_false/input_0.cairo b/tests/nodes/argmax_u32_2D_keepdims_false/input_0.cairo
deleted file mode 100644
index 19a457838..000000000
--- a/tests/nodes/argmax_u32_2D_keepdims_false/input_0.cairo
+++ /dev/null
@@ -1,16 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn input_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(207);
-    data.append(36);
-    data.append(209);
-    data.append(4);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_2D_keepdims_false/output_0.cairo b/tests/nodes/argmax_u32_2D_keepdims_false/output_0.cairo
deleted file mode 100644
index edc68637b..000000000
--- a/tests/nodes/argmax_u32_2D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,13 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_2D_last_index.cairo b/tests/nodes/argmax_u32_2D_last_index.cairo
deleted file mode 100644
index 60b70f828..000000000
--- a/tests/nodes/argmax_u32_2D_last_index.cairo
+++ /dev/null
@@ -1,20 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_u32_2D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_u32_2D_last_index/input_0.cairo b/tests/nodes/argmax_u32_2D_last_index/input_0.cairo
deleted file mode 100644
index 0f577ecf4..000000000
--- a/tests/nodes/argmax_u32_2D_last_index/input_0.cairo
+++ /dev/null
@@ -1,16 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn input_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(41);
-    data.append(93);
-    data.append(233);
-    data.append(71);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_2D_last_index/output_0.cairo b/tests/nodes/argmax_u32_2D_last_index/output_0.cairo
deleted file mode 100644
index 221e8a246..000000000
--- a/tests/nodes/argmax_u32_2D_last_index/output_0.cairo
+++ /dev/null
@@ -1,14 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_3D_default.cairo b/tests/nodes/argmax_u32_3D_default.cairo
deleted file mode 100644
index f9f47d2fd..000000000
--- a/tests/nodes/argmax_u32_3D_default.cairo
+++ /dev/null
@@ -1,20 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_u32_3D_default() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_u32_3D_default/input_0.cairo b/tests/nodes/argmax_u32_3D_default/input_0.cairo
deleted file mode 100644
index 83343807d..000000000
--- a/tests/nodes/argmax_u32_3D_default/input_0.cairo
+++ /dev/null
@@ -1,21 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn input_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(22);
-    data.append(254);
-    data.append(48);
-    data.append(151);
-    data.append(21);
-    data.append(13);
-    data.append(254);
-    data.append(100);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_3D_default/output_0.cairo b/tests/nodes/argmax_u32_3D_default/output_0.cairo
deleted file mode 100644
index e3da7dd7d..000000000
--- a/tests/nodes/argmax_u32_3D_default/output_0.cairo
+++ /dev/null
@@ -1,17 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(0);
-    data.append(0);
-    data.append(1);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_3D_keepdims_false.cairo b/tests/nodes/argmax_u32_3D_keepdims_false.cairo
deleted file mode 100644
index 7ee7d2aad..000000000
--- a/tests/nodes/argmax_u32_3D_keepdims_false.cairo
+++ /dev/null
@@ -1,20 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_u32_3D_keepdims_false() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::Some(false), Option::None(()));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_u32_3D_keepdims_false/input_0.cairo b/tests/nodes/argmax_u32_3D_keepdims_false/input_0.cairo
deleted file mode 100644
index 4842ff00c..000000000
--- a/tests/nodes/argmax_u32_3D_keepdims_false/input_0.cairo
+++ /dev/null
@@ -1,21 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn input_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(183);
-    data.append(138);
-    data.append(72);
-    data.append(104);
-    data.append(239);
-    data.append(185);
-    data.append(16);
-    data.append(188);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_3D_keepdims_false/output_0.cairo b/tests/nodes/argmax_u32_3D_keepdims_false/output_0.cairo
deleted file mode 100644
index 15d30cb7a..000000000
--- a/tests/nodes/argmax_u32_3D_keepdims_false/output_0.cairo
+++ /dev/null
@@ -1,16 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    data.append(1);
-    data.append(0);
-    data.append(1);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_3D_last_index.cairo b/tests/nodes/argmax_u32_3D_last_index.cairo
deleted file mode 100644
index 8098105e1..000000000
--- a/tests/nodes/argmax_u32_3D_last_index.cairo
+++ /dev/null
@@ -1,20 +0,0 @@
-mod input_0;
-mod output_0;
-
-
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::U32TensorPartialEq;
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-use orion::utils::{assert_eq, assert_seq_eq};
-use orion::operators::tensor::{TensorTrait, Tensor};
-
-#[test]
-#[available_gas(2000000000)]
-fn test_argmax_u32_3D_last_index() {
-    let input_0 = input_0::input_0();
-    let z = output_0::output_0();
-
-    let y = input_0.argmax(0, Option::None(()), Option::Some(true));
-
-    assert_eq(y, z);
-}
diff --git a/tests/nodes/argmax_u32_3D_last_index/input_0.cairo b/tests/nodes/argmax_u32_3D_last_index/input_0.cairo
deleted file mode 100644
index 2a464efd3..000000000
--- a/tests/nodes/argmax_u32_3D_last_index/input_0.cairo
+++ /dev/null
@@ -1,21 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn input_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(2);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(9);
-    data.append(169);
-    data.append(140);
-    data.append(99);
-    data.append(130);
-    data.append(132);
-    data.append(79);
-    data.append(57);
-    TensorTrait::new(shape.span(), data.span())
-}
diff --git a/tests/nodes/argmax_u32_3D_last_index/output_0.cairo b/tests/nodes/argmax_u32_3D_last_index/output_0.cairo
deleted file mode 100644
index 217f9ee86..000000000
--- a/tests/nodes/argmax_u32_3D_last_index/output_0.cairo
+++ /dev/null
@@ -1,17 +0,0 @@
-use core::array::{ArrayTrait, SpanTrait};
-use orion::operators::tensor::{TensorTrait, Tensor};
-use orion::operators::tensor::{U32Tensor, U32TensorAdd};
-
-fn output_0() -> Tensor<u32> {
-    let mut shape = ArrayTrait::<usize>::new();
-    shape.append(1);
-    shape.append(2);
-    shape.append(2);
-
-    let mut data = ArrayTrait::new();
-    data.append(1);
-    data.append(0);
-    data.append(0);
-    data.append(0);
-    TensorTrait::new(shape.span(), data.span())
-}