diff --git a/include/aidge/backend/cpu.hpp b/include/aidge/backend/cpu.hpp
index b81145dec13d39ea701f2ae9dfcbd6b6d17bff39..d9e27d96b4c60d947ddc6a965276531aa7907d89 100644
--- a/include/aidge/backend/cpu.hpp
+++ b/include/aidge/backend/cpu.hpp
@@ -13,6 +13,7 @@
#define AIDGE_CPU_IMPORTS_H_
#include "aidge/backend/cpu/operator/AddImpl.hpp"
+#include "aidge/backend/cpu/operator/ArgMaxImpl.hpp"
#include "aidge/backend/cpu/operator/AvgPoolingImpl.hpp"
#include "aidge/backend/cpu/operator/MaxPoolingImpl.hpp"
#include "aidge/backend/cpu/operator/BatchNormImpl.hpp"
diff --git a/include/aidge/backend/cpu/operator/ArgMaxImpl.hpp b/include/aidge/backend/cpu/operator/ArgMaxImpl.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..f93abbbca9630a8b60290bae7660f6428b41b0b3
--- /dev/null
+++ b/include/aidge/backend/cpu/operator/ArgMaxImpl.hpp
@@ -0,0 +1,60 @@
+/********************************************************************************
+ * Copyright (c) 2024 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#ifndef AIDGE_CPU_OPERATOR_ARGMAXIMPL_H_
+#define AIDGE_CPU_OPERATOR_ARGMAXIMPL_H_
+
+#include <array>
+#include <memory>
+#include <tuple>
+#include <vector>
+
+#include "aidge/backend/OperatorImpl.hpp"
+#include "aidge/operator/ArgMax.hpp"
+#include "aidge/utils/Registrar.hpp"
+#include "aidge/utils/Types.h"
+
+namespace Aidge {
+class ArgMaxImplForward_cpu
+ : public Registrable<ArgMaxImplForward_cpu,
+ std::tuple<DataType, DataType>,
+ void(std::int32_t,
+ DimSize_t,
+ const std::vector<DimSize_t>&,
+ const void *,
+ void *)> {};
+class ArgMaxImplBackward_cpu
+ : public Registrable<ArgMaxImplBackward_cpu,
+ std::tuple<DataType, DataType>,
+ void(std::int32_t,
+ DimSize_t,
+ const std::vector<DimSize_t>&,
+ const void *,
+ void *)> {};
+
+class ArgMaxImpl_cpu : public OperatorImpl {
+ public:
+ ArgMaxImpl_cpu(const ArgMax_Op& op) : OperatorImpl(op, "cpu") {}
+
+ static std::unique_ptr<ArgMaxImpl_cpu> create(const ArgMax_Op &op) {
+ return std::make_unique<ArgMaxImpl_cpu>(op);
+ }
+
+ public:
+ void forward() override;
+};
+
+namespace {
+static Registrar<ArgMax_Op> registrarArgMaxImpl_cpu("cpu", Aidge::ArgMaxImpl_cpu::create);
+} // namespace
+} // namespace Aidge
+
+#endif /* AIDGE_CPU_OPERATOR_ARGMAXIMPL_H_ */
diff --git a/include/aidge/backend/cpu/operator/ArgMaxImpl_forward_kernels.hpp b/include/aidge/backend/cpu/operator/ArgMaxImpl_forward_kernels.hpp
new file mode 100644
index 0000000000000000000000000000000000000000..a03a02441ebbc81f923b4506b9478f89e461b01b
--- /dev/null
+++ b/include/aidge/backend/cpu/operator/ArgMaxImpl_forward_kernels.hpp
@@ -0,0 +1,90 @@
+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#ifndef AIDGE_CPU_OPERATOR_ARGMAXIMPL_FORWARD_KERNEL_H_
+#define AIDGE_CPU_OPERATOR_ARGMAXIMPL_FORWARD_KERNEL_H_
+
+#include <algorithm> // std::for_each
+#include <cstddef> // std::size_t
+#include <cstdint> // std::int32_t
+#include <functional> //std::multiplies
+#include <numeric> //std::accumulate
+#include <vector>
+#include <limits>
+
+#include "aidge/backend/cpu/operator/ArgMaxImpl.hpp"
+#include "aidge/data/Data.hpp"
+#include "aidge/operator/ArgMax.hpp"
+#include "aidge/utils/Registrar.hpp"
+
+namespace Aidge {
+template <class I, class O>
+void ArgMaxImpl_cpu_forward_kernel(std::int32_t axis_,
+ DimSize_t select_last_index,
+ const std::vector<DimSize_t>& inputDims,
+ const void* input_,
+ void* output_) {
+
+ const I* input = static_cast<const I*>(input_);
+ O* output = static_cast<O*>(output_);
+
+ const std::size_t axis = static_cast<std::size_t>(axis_);
+
+ const std::size_t nb_dims = inputDims.size();
+
+ auto stride_post = std::unique_ptr<std::size_t[]>(new std::size_t[nb_dims]);
+ stride_post[nb_dims - 1] = 1;
+ for (std::size_t i = nb_dims-2; i != static_cast<std::size_t>(-1); --i) {
+ stride_post[i] = stride_post[i+1]*inputDims[i+1];
+ }
+ auto stride_pre = std::unique_ptr<std::size_t[]>(new std::size_t[nb_dims]);
+ stride_pre[0] = 1;
+ for (std::size_t i = 1; i < nb_dims; ++i) {
+ stride_pre[i] = stride_pre[i-1]*inputDims[i-1];
+ }
+
+ const std::size_t dim_i = inputDims[axis];
+ for (std::size_t pre = 0; pre < stride_pre[axis]; ++pre) {
+ for (std::size_t post = 0; post < stride_post[axis]; ++post) {
+ const std::size_t idx_i = pre * dim_i * stride_post[axis] + post;
+ const std::size_t idx_o = pre * stride_post[axis] + post;
+ I max = std::numeric_limits<I>::min();
+ for (std::size_t i = 0; i < dim_i; ++i) {
+ if (select_last_index) {
+ if (input[idx_i]>=max)
+ {
+ output[idx_o] = i;
+ }
+ }
+ else {
+ if (input[idx_i] > max)
+ {
+ output[idx_o] = i;
+ }
+ }
+
+ }
+ }
+ }
+
+}
+
+namespace {
+static Registrar<ArgMaxImplForward_cpu> registrarArgMaxImplForward_cpu_Float32(
+ {DataType::Float32, DataType::Float32}, Aidge::ArgMaxImpl_cpu_forward_kernel<float, float>);
+static Registrar<ArgMaxImplForward_cpu> registrarArgMaxImplForward_cpu_Int32(
+ {DataType::Int32, DataType::Int32}, Aidge::ArgMaxImpl_cpu_forward_kernel<int, int>);
+static Registrar<ArgMaxImplForward_cpu> registrarArgMaxImplForward_cpu_Float64(
+ {DataType::Float64, DataType::Float64}, Aidge::ArgMaxImpl_cpu_forward_kernel<double, double>);
+} // namespace
+} // namespace Aidge
+
+#endif /* AIDGE_CPU_OPERATOR_ARGMAXIMPL_FORWARD_KERNEL_H_ */
diff --git a/src/operator/ArgMaxImpl.cpp b/src/operator/ArgMaxImpl.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..eda3c0b2de62e8e170c2562945999273bdb921a7
--- /dev/null
+++ b/src/operator/ArgMaxImpl.cpp
@@ -0,0 +1,34 @@
+/********************************************************************************
+ * Copyright (c) 2024 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#include "aidge/backend/cpu/operator/ArgMaxImpl.hpp"
+
+#include <memory>
+#include <vector>
+
+#include "aidge/utils/Types.h"
+#include "aidge/operator/ArgMax.hpp"
+#include "aidge/backend/cpu/operator/ArgMaxImpl_forward_kernels.hpp"
+
+void Aidge::ArgMaxImpl_cpu::forward() {
+ const ArgMax_Op& op_ = dynamic_cast<const ArgMax_Op&>(mOp);
+ // Find the correct kernel type
+ auto kernelFunc = Registrar<ArgMaxImplForward_cpu>::create({
+ op_.getInput(0)->dataType(),
+ op_.getOutput(0)->dataType()});
+
+ // Call kernel
+ kernelFunc(op_.axis(),
+ op_.selectLastIndex(),
+ op_.getInput(0)->dims(),
+ op_.getInput(0)->getImpl()->rawPtr(),
+ op_.getOutput(0)->getImpl()->rawPtr());
+}
diff --git a/unit_tests/operator/Test_ArgMax.cpp b/unit_tests/operator/Test_ArgMax.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..1851d929f39f246a5331da8bbea7119bc2e3875a
--- /dev/null
+++ b/unit_tests/operator/Test_ArgMax.cpp
@@ -0,0 +1,116 @@
+/********************************************************************************
+ * Copyright (c) 2024 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#include <catch2/catch_test_macros.hpp>
+#include <memory>
+
+#include "aidge/data/Tensor.hpp"
+#include "aidge/operator/ArgMax.hpp"
+#include "aidge/operator/Conv.hpp"
+
+#include "aidge/backend/cpu.hpp"
+#include "aidge/utils/TensorUtils.hpp"
+
+using namespace Aidge;
+
+TEST_CASE("[cpu/operator] ArgMax(forward)", "[ArgMax][CPU]") {
+ SECTION("3D Tensor") {
+ std::shared_ptr<Tensor> myInput = std::make_shared<Tensor>(Array3D<float,2,3,4> {
+ {
+ {
+ { 1.0, 2.0, 3.0, 4.0},
+ { 8.0, 0.0, 17.0, 1.0},
+ { 5.0, 10.0, 6.0, 0.0}
+ },
+ {
+ { 7.0, 1.0, 9.0, 4.0},
+ { 0.0, 8.0, 4.0, 2.0},
+ { 9.0, 2.0, 0.0, 5.0}
+ }
+ }
+ });
+ SECTION("Axis 2") {
+
+ Tensor myOutput = Tensor(Array2D<float,2,3> {
+ {
+ { 3.0, 2.0, 1.0 },
+ { 2.0, 1.0, 0.0}
+ }
+ });
+
+ std::shared_ptr<Node> myArgMax = ArgMax(2);
+ auto op = std::static_pointer_cast<OperatorTensor>(myArgMax -> getOperator());
+ op->associateInput(0,myInput);
+ op->setDataType(DataType::Float32);
+ op->setBackend("cpu");
+ myArgMax->forward();
+ op->getOutput(0)->print();
+
+ REQUIRE(*(op->getOutput(0)) == myOutput);
+ }
+ SECTION("Axis 1") {
+ Tensor myOutput = Tensor(Array2D<float,2,4> {
+ {
+ { 1.0, 2.0, 1.0, 0.0 },
+ { 2.0, 1.0, 0.0, 2.0 }
+ }
+ });
+
+ std::shared_ptr<Node> myArgMax = ArgMax(1);
+ auto op = std::static_pointer_cast<OperatorTensor>(myArgMax -> getOperator());
+ op->associateInput(0,myInput);
+ op->setDataType(DataType::Float32);
+ op->setBackend("cpu");
+ myArgMax->forward();
+ myOutput.print();
+ op->getOutput(0)->print();
+ REQUIRE(*(op->getOutput(0)) == myOutput);
+ }
+ SECTION("Axis 0") {
+ Tensor myOutput = Tensor(Array2D<float,3,4> {
+ {
+ { 1.0, 0.0, 1.0, 0.0 },
+ { 0.0, 1.0, 0.0, 1.0 },
+ { 1.0, 0.0, 0.0, 1.0 }
+ },
+ });
+
+ std::shared_ptr<Node> myArgMax = ArgMax(1);
+ auto op = std::static_pointer_cast<OperatorTensor>(myArgMax -> getOperator());
+ op->associateInput(0,myInput);
+ op->setDataType(DataType::Float32);
+ op->setBackend("cpu");
+ myArgMax->forward();
+ myOutput.print();
+ op->getOutput(0)->print();
+ REQUIRE(*(op->getOutput(0)) == myOutput);
+ }
+ }
+ SECTION("Select_Last_Index") {
+ std::shared_ptr<Tensor> myInput = std::make_shared<Tensor>(Array1D<float,10> {
+ {
+ 1.0, 5.0, 9.0, 0.0, 6.0, 2.0, 9.0, 4.0, 3.0, 9.0
+ }
+ });
+ std::shared_ptr<Tensor> myOutput = std::make_shared<Tensor>(Array1D<float,1> {{9}});
+
+ std::shared_ptr<Node> myArgMax = ArgMax(0, 1, 1);
+ auto op = std::static_pointer_cast<OperatorTensor>(myArgMax -> getOperator());
+ op->associateInput(0,myInput);
+ op->setDataType(DataType::Float32);
+ op->setBackend("cpu");
+ myArgMax->forward();
+ op->getOutput(0)->print();
+
+ REQUIRE(*(op->getOutput(0)) == *myOutput);
+
+ }
+}
\ No newline at end of file