From c90fcc93cb1fc008b7ed7065a1961fd10bbb1525 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gr=C3=A9goire=20KUBLER?= <gregoire.kubler@proton.me>
Date: Fri, 6 Sep 2024 15:43:46 +0200
Subject: [PATCH] feat : unsqueeze operator

---
 .../unit_tests/test_operator_unsqueeze.py     | 211 ++++++++++
 include/aidge/operator/Unsqueeze.hpp          | 157 +++++++
 python_binding/operator/pybind_Unsqueeze.cpp  |  48 +++
 python_binding/pybind_core.cpp                |   2 +
 src/operator/Unsqueeze.cpp                    | 127 ++++++
 unit_tests/operator/Test_Unsqueeze_Op.cpp     | 382 ++++++++++++++++++
 6 files changed, 927 insertions(+)
 create mode 100644 aidge_core/unit_tests/test_operator_unsqueeze.py
 create mode 100644 include/aidge/operator/Unsqueeze.hpp
 create mode 100644 python_binding/operator/pybind_Unsqueeze.cpp
 create mode 100644 src/operator/Unsqueeze.cpp
 create mode 100644 unit_tests/operator/Test_Unsqueeze_Op.cpp

diff --git a/aidge_core/unit_tests/test_operator_unsqueeze.py b/aidge_core/unit_tests/test_operator_unsqueeze.py
new file mode 100644
index 000000000..12f55fa30
--- /dev/null
+++ b/aidge_core/unit_tests/test_operator_unsqueeze.py
@@ -0,0 +1,211 @@
+"""
+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
+"""
+
+import unittest
+import aidge_core
+import numpy as np
+from numpy import testing as npt
+
+
+class TestUnsqueeze(unittest.TestCase):
+    """
+    Test unsqueeze operator
+    """
+
+    def setUp(self):
+        axis_to_unsqueeze_dim_0 = [0]
+        axis_to_unsqueeze_many = [1, 4, 5]
+        axis_to_unsqueeze_error_identical_index = [0, 0, 0]
+        axis_to_unsqueeze_error_too_high_index = [50]
+        axis_to_unsqueeze_onnx_test = [0, 4]
+        unsqueeze_dim_0 = aidge_core.Unsqueeze(
+            axis_to_unsqueeze_dim_0, name="unsqueeze_dim_0"
+        )
+        unsqueeze_many = aidge_core.Unsqueeze(
+            axis_to_unsqueeze_many, name="unsqueeze_many"
+        )
+        unsqueeze_error_identical_index = aidge_core.Unsqueeze(
+            axis_to_unsqueeze_error_identical_index,
+            name="unsqueeze_error_identical_index",
+        )
+        unsqueeze_error_node = aidge_core.Unsqueeze(
+            axis_to_unsqueeze_error_too_high_index,
+            name="unsqueeze_error_index_too_high",
+        )
+        unsqueeze_onnx_test = aidge_core.Unsqueeze(
+            axis_to_unsqueeze_onnx_test, name="unsqueeze taken from onnx documentation"
+        )
+
+        input_1_data_shape = np.array([1, 2, 3])
+        input_2_data_shape = np.array([2, 1, 3, 3])
+        input_3_data_shape = np.array([1, 1, 4])
+        input_onnx_data_shape = np.array([3, 4, 5])
+
+        input_axes_dim_0 = axis_to_unsqueeze_dim_0
+        input_axes_many = axis_to_unsqueeze_many
+        input_axes_onnx_test = axis_to_unsqueeze_onnx_test
+
+        self.tests_axes_defined_by_attribute = [
+            (input_1_data_shape, unsqueeze_dim_0, np.array([1, 1, 2, 3])),
+            (input_2_data_shape, unsqueeze_dim_0, np.array([1, 2, 1, 3, 3])),
+            (input_2_data_shape, unsqueeze_many, np.array([2, 1, 1, 3, 1, 1, 3])),
+            (input_3_data_shape, unsqueeze_dim_0, np.array([1, 1, 1, 4])),
+            (input_3_data_shape, unsqueeze_many, np.array([1, 1, 1, 4, 1, 1])),
+            (input_onnx_data_shape, unsqueeze_onnx_test, np.array([1, 3, 4, 5, 1])),
+        ]
+
+        self.tests_axes_defined_by_tensor = [
+            (
+                input_1_data_shape,
+                input_axes_dim_0,
+                unsqueeze_error_node,
+                np.array([1, 1, 2, 3]),
+            ),
+            (
+                input_2_data_shape,
+                input_axes_dim_0,
+                unsqueeze_error_node,
+                np.array([1, 2, 1, 3, 3]),
+            ),
+            (
+                input_2_data_shape,
+                input_axes_many,
+                unsqueeze_error_node,
+                np.array([2, 1, 1, 3, 1, 1, 3]),
+            ),
+            (
+                input_3_data_shape,
+                input_axes_dim_0,
+                unsqueeze_error_node,
+                np.array([1, 1, 1, 4]),
+            ),
+            (
+                input_3_data_shape,
+                input_axes_many,
+                unsqueeze_error_node,
+                np.array([1, 1, 1, 4, 1, 1]),
+            ),
+            (
+                input_onnx_data_shape,
+                input_axes_onnx_test,
+                unsqueeze_error_node,
+                np.array([1, 3, 4, 5, 1]),
+            ),
+        ]
+
+        self.test_error = [
+            (input_1_data_shape, unsqueeze_error_identical_index),
+            (input_1_data_shape, unsqueeze_error_node),
+            (input_1_data_shape, unsqueeze_many),  # dims too high
+        ]
+        return
+
+    def tearDown(self):
+        pass
+
+    def test_axes_defined_by_attribute(self):
+        for index, (
+            input_shape,
+            unsqueeze_template,
+            expected_output_shape,
+        ) in enumerate(self.tests_axes_defined_by_attribute):
+            test_unsqueeze = unsqueeze_template
+            test_unsqueeze_op = test_unsqueeze.get_operator()
+
+            print(f"\nTest {index}")
+            print(f"input size : {input_shape}")
+            print(f"operator : {test_unsqueeze}")
+            print(f"expected output_shape : {expected_output_shape}")
+
+            test_unsqueeze_op.set_backend("cpu")
+
+            input_values = np.ones(shape=input_shape, dtype=np.float32)
+            expected_output_values = np.ones(
+                shape=expected_output_shape, dtype=np.float32
+            )
+            input_tensor = aidge_core.Tensor(input_values)
+            test_unsqueeze_op.set_input(0, input_tensor)
+
+            test_unsqueeze_op.forward_dims()
+            test_unsqueeze_op.forward()
+
+            unsqueeze_output = test_unsqueeze_op.get_output(0)
+
+            npt.assert_array_equal(
+                unsqueeze_output.dims(),
+                expected_output_shape,
+                err_msg=f"UNSQUEEZE FAILURE : expected result dimensions differs from output's\n\toperator : {test_unsqueeze}\n\tinput.shape : {input_shape.shape}",
+            )
+            npt.assert_array_almost_equal(
+                np.array(unsqueeze_output),
+                expected_output_values,
+                7,
+                err_msg=f"UNSQUEEZE FAILURE : output tensor values differs from expected values\n\toperator : {test_unsqueeze}\n\tinput.shape : {input_shape.shape}",
+            )
+        return
+
+    def test_axes_defined_via_tensor_input(self):
+        for index, (
+            input_shape,
+            input_axes_to_squeeze,
+            squeeze_node_template,
+            output_shape,
+        ) in enumerate(self.tests_axes_defined_by_tensor):
+            test_squeeze_node = squeeze_node_template
+            test_squeeze_op = test_squeeze_node.get_operator()
+
+            print(f"\nTest {index}")
+            print(f"input shape : {input_shape}")
+            print(f"input axes: {np.array(input_axes_to_squeeze)}")
+            print(f"operator : {test_squeeze_node}")
+            print(f"expected output_shape : {output_shape}")
+
+            test_squeeze_op.set_backend("cpu")
+            test_squeeze_op.set_datatype(aidge_core.dtype.float32)
+
+            input_values = np.ones(shape=input_shape, dtype=np.float32)
+            output_values = np.ones(shape=output_shape, dtype=np.float32)
+
+            input_data = aidge_core.Tensor(input_values)
+            input_data.set_datatype(aidge_core.dtype.float32)
+            input_data.set_backend("cpu")
+
+            input_axes = aidge_core.Tensor(
+                np.array(input_axes_to_squeeze, dtype=np.float32)
+            )
+            input_axes.set_datatype(aidge_core.dtype.int8)
+            input_axes.set_backend("cpu")
+
+            test_squeeze_op.set_input(0, input_data)
+            test_squeeze_op.set_input(1, input_axes)
+
+            self.assertEqual(test_squeeze_op.forward_dims(True), True)
+            test_squeeze_op.forward()
+
+            squeeze_output = test_squeeze_op.get_output(0)
+
+            npt.assert_array_equal(
+                squeeze_output.dims(),
+                output_shape,
+                err_msg=f"SQUEEZE FAILURE : expected result differs from output size\n\toperator : {test_squeeze_node}\n\tinput.shape : {input_shape.shape}",
+            )
+            npt.assert_array_almost_equal(
+                np.array(squeeze_output, dtype=np.float32),
+                output_values,
+                7,
+                err_msg=f"SQUEEZE FAILURE : output tensor values differs from expected values\n\toperator : {test_squeeze_node}\n\tinput.shape : {input_shape.shape}",
+            )
+            # self.assertEqual(test_squeeze_op.dims_forwarded(), True, "SQUEEZE_FAILURE : dims_forwarded failed.")
+        return
+
+
+if __name__ == "__main__":
+    unittest.main()
+
diff --git a/include/aidge/operator/Unsqueeze.hpp b/include/aidge/operator/Unsqueeze.hpp
new file mode 100644
index 000000000..3443801bc
--- /dev/null
+++ b/include/aidge/operator/Unsqueeze.hpp
@@ -0,0 +1,157 @@
+/********************************************************************************
+ * 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_CORE_OPERATOR_UNSQUEEZE_H_
+#define AIDGE_CORE_OPERATOR_UNSQUEEZE_H_
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "aidge/graph/Node.hpp"
+#include "aidge/operator/Operator.hpp"
+#include "aidge/operator/OperatorTensor.hpp"
+#include "aidge/utils/ErrorHandling.hpp"
+#include "aidge/utils/Registrar.hpp"
+#include "aidge/utils/StaticAttributes.hpp"
+#include "aidge/utils/Types.h"
+
+namespace Aidge {
+/**
+ * @brief implementation of the operator unsqueeze.
+ * @note Since this operator implementation is agnostic to the backend it is
+ * located here instead of in aidge_backend_cpu/cuda.
+ */
+class Unsqueeze_OpImpl : public OperatorImpl {
+public:
+  Unsqueeze_OpImpl(const Operator &op, const std::string &backend = "")
+      : OperatorImpl(op, backend) {}
+  void forward() override;
+};
+
+enum class UnsqueezeAttr {
+  /**
+   * @brief vector of axes to unsqueeze.
+   * values must be comprised within
+   * [ -a ; a-1 ]
+   * with a = input_tensor.nbDim() + dims_to_unsqueeze.size()
+   */
+  Axes
+};
+
+/**
+ * @brief This operator has as purpose to add a dummy dimension around given
+ * axis. Unsqueezing the 2nd dim of a tensor of dim (1,2,3,4) will result in a
+ * tensor of dim (1,2,1,3,4)
+ * You can also unsqueeze dimensions whose index is higher than the nb of input
+ * dimensions as long as :
+ * dims_to_unsqueeze[i] < tensor.nbDim() +
+ * dims_to_unsqueeze.size()
+ */
+class Unsqueeze_Op
+    : public OperatorTensor,
+      public Registrable<Unsqueeze_Op, std::string,
+                         std::shared_ptr<OperatorImpl>(const Unsqueeze_Op &)> {
+
+public:
+  static const std::string
+      Type; // name of the type of the operation (Here "Unsqueeze")
+
+private:
+  using Attributes_ = StaticAttributes<UnsqueezeAttr, std::vector<int8_t>>;
+  template <UnsqueezeAttr e>
+  using attr = typename Attributes_::template attr<e>;
+  const std::shared_ptr<Attributes_> mAttributes;
+
+public:
+  Unsqueeze_Op() =
+      delete; // no default constructor since this class has attributes
+
+  /**
+   * @brief constructor for Unsqueeze op
+   * @param[in] axis around which perform the operation
+   */
+  Unsqueeze_Op(const std::vector<int8_t> &axes)
+      : OperatorTensor(Type, {InputCategory::Data, InputCategory::OptionalData},
+                       1),
+        mAttributes(
+            std::make_shared<Attributes_>(attr<UnsqueezeAttr::Axes>(axes))) {
+    mImpl = std::make_shared<Unsqueeze_OpImpl>(*this);
+  }
+
+  /**
+   * @brief Copy-constructor. Copy the operator attributes and its output
+   * tensor(s), but not its input tensors (the new operator has no input
+   * associated).
+   * @param op Operator to copy.
+   */
+  Unsqueeze_Op(const Unsqueeze_Op &op)
+      : OperatorTensor(op), mAttributes(op.mAttributes) {
+    if (!op.backend().empty()) {
+      SET_IMPL_MACRO(Unsqueeze_Op, *this, op.backend());
+    } else {
+      mImpl = std::make_shared<Unsqueeze_OpImpl>(*this);
+    }
+  }
+
+  /**
+   * @brief Clone the operator using its copy-constructor.
+   * @see Operator::MatMul_Op
+   */
+  std::shared_ptr<Operator> clone() const override final {
+    return std::make_shared<Unsqueeze_Op>(*this);
+  }
+
+  /**
+   * @brief Compute dimensions for the output Tensor
+   */
+  bool forwardDims(bool allowDataDependency = false) override final;
+  bool dimsForwarded() const override final;
+
+  void setBackend(const std::string &name,
+                  DeviceIdx_t device = 0) override final;
+
+  inline std::shared_ptr<Attributes> attributes() const override {
+    return mAttributes;
+  }
+  /**
+   * @brief vector of axes to unsqueeze.
+   * values must be comprised within
+   * [ -a ; a-1 ]
+   * with : a = input_tensor.nbDim() + dims_to_unsqueeze.size()
+   */
+  inline std::vector<int8_t> &axes() const noexcept {
+    return mAttributes->template getAttr<UnsqueezeAttr::Axes>();
+  }
+
+  static const std::vector<std::string> getInputsName() {
+    return {"data_input", "axes_to_unsqueeze"};
+  }
+  static const std::vector<std::string> getOutputsName() {
+    return {"unsqueezed"};
+  }
+};
+
+// helper with C-style array instead of std::array for kernel_dims to allow
+// automatic template DIM deduction
+inline std::shared_ptr<Node> Unsqueeze(const std::vector<int8_t> &axes = {},
+                                       const std::string &name = "") {
+  return std::make_shared<Node>(std::make_shared<Unsqueeze_Op>(axes), name);
+}
+} // namespace Aidge
+
+namespace {
+template <>
+const char *const EnumStrings<Aidge::UnsqueezeAttr>::data[] = {"Axes"};
+}
+
+#endif // AIDGE_CORE_OPERATOR_UNSQUEEZE_H_
diff --git a/python_binding/operator/pybind_Unsqueeze.cpp b/python_binding/operator/pybind_Unsqueeze.cpp
new file mode 100644
index 000000000..40c179c40
--- /dev/null
+++ b/python_binding/operator/pybind_Unsqueeze.cpp
@@ -0,0 +1,48 @@
+/********************************************************************************
+ * 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
+ *
+ ********************************************************************************/
+
+#include <pybind11/pybind11.h>
+#include <string>
+#include <vector>
+
+#include "aidge/operator/OperatorTensor.hpp"
+#include "aidge/operator/Unsqueeze.hpp"
+
+namespace py = pybind11;
+namespace Aidge {
+
+void init_Unsqueeze(py::module &m) {
+  py::class_<Unsqueeze_Op, std::shared_ptr<Unsqueeze_Op>, OperatorTensor>(
+      m, "UnsqueezeOp", py::multiple_inheritance(),
+      R"mydelimiter(
+		Initialize an unsqueeze operator.
+		:param axes :   axes to unsqueeze between [-r;r-1] 
+						with r = input_tensor.nbDims() + len(axes)
+		:type axes : :py:class: List[Int]
+		)mydelimiter")
+      // Here we bind the methods of the Unsqueeze_Op that wil want to access
+      .def("get_inputs_name", &Unsqueeze_Op::getInputsName)
+      .def("get_outputs_name", &Unsqueeze_Op::getOutputsName)
+      .def("axes", &Unsqueeze_Op::axes);
+  // Here we bind the constructor of the Unsqueeze Node. We add an argument for
+  // each attribute of the operator (in here we only have 'axes') and the last
+  // argument is the node's name.
+  m.def("Unsqueeze", &Unsqueeze, py::arg("axes") = std::vector<int8_t>({}),
+        py::arg("name") = "",
+        R"mydelimiter(
+    Initialize a node containing an unsqueeze operator.
+	:param axes :   axes to unsqueeze between [-r;r-1] 
+					with r = input_tensor.nbDims() + len(axes)
+	:type axes : :py:class: List[Int]
+    :param name : name of the node.
+)mydelimiter");
+}
+} // namespace Aidge
diff --git a/python_binding/pybind_core.cpp b/python_binding/pybind_core.cpp
index cc23727be..8f0131a37 100644
--- a/python_binding/pybind_core.cpp
+++ b/python_binding/pybind_core.cpp
@@ -68,6 +68,7 @@ void init_Sub(py::module&);
 void init_Tanh(py::module&);
 void init_Transpose(py::module&);
 void init_Identity(py::module&);
+void init_Unsqueeze(py::module&);
 
 void init_Node(py::module&);
 void init_GraphView(py::module&);
@@ -144,6 +145,7 @@ void init_Aidge(py::module& m) {
     init_Tanh(m);
     init_Transpose(m);
     init_Identity(m);
+    init_Unsqueeze(m);
 
     init_Producer(m);
 
diff --git a/src/operator/Unsqueeze.cpp b/src/operator/Unsqueeze.cpp
new file mode 100644
index 000000000..e88e0f8ca
--- /dev/null
+++ b/src/operator/Unsqueeze.cpp
@@ -0,0 +1,127 @@
+/********************************************************************************
+ * 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
+ *
+ ********************************************************************************/
+
+#include "aidge/operator/Unsqueeze.hpp"
+
+#include <cstdint>
+#include <fmt/core.h>
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "aidge/data/Data.hpp"
+#include "aidge/data/Tensor.hpp"
+#include "aidge/utils/ErrorHandling.hpp"
+#include "aidge/utils/Log.hpp"
+#include "aidge/utils/Registrar.hpp"
+#include "aidge/utils/Types.h"
+
+namespace Aidge {
+const std::string Unsqueeze_Op::Type = "Unsqueeze";
+
+bool Aidge::Unsqueeze_Op::dimsForwarded() const {
+  if ((getInput(1) && !getInput(1)->undefined())) {
+    // output dims are data dependent
+    return false;
+  }
+
+  return OperatorTensor::dimsForwarded();
+}
+
+bool Unsqueeze_Op::forwardDims(bool allowDataDependency) {
+  // error checking
+  if (!inputsAssociated(true)) {
+    return false;
+  }
+  std::shared_ptr<Tensor> fallback;
+  // Copy optional input #1, if present, to attribute Axes
+  if (getInput(1)) {
+    if (!this->axes().empty()) {
+      Log::notice("{} : ignoring non-empty \"axes\" attribute because input#1 "
+                  "takes precedence",
+                  type());
+    }
+
+    if (!allowDataDependency) {
+      Log::warn("{} : unable to forwardDims() because output dims are data "
+                "dependent on input#1",
+                type());
+      return false;
+    }
+
+    this->axes().clear(); // If both are provided input would override attrs
+    this->axes().reserve(getInput(1)->size());
+    const auto &axes =
+        getInput(1)->refCastFrom(fallback, NativeType<int8_t>::type, "cpu");
+    std::copy_n(static_cast<int8_t *>(axes.getImpl()->hostPtr()),
+                axes.size(), std::back_inserter(this->axes()));
+  }
+  AIDGE_ASSERT(!this->axes().empty(),
+               "{} : Axes to unsqueeze can be defined via input#1 or axes "
+               "attribute. None of them were provided.",
+               type());
+
+  std::vector<DimSize_t> input_dims = getInput(0)->dims();
+  std::vector<DimIdx_t> axes_rectified_idx;
+  axes_rectified_idx.reserve(this->axes().size());
+  DimIdx_t output_nb_dims = input_dims.size() + this->axes().size();
+
+  for (const int8_t &axis : this->axes()) {
+    AIDGE_ASSERT(axis >= static_cast<int8_t>(-output_nb_dims) &&
+                     axis < static_cast<int8_t>(output_nb_dims),
+                 "{} : Axis index OutOfBounds enrror, expected value "
+                 "within size limits of input tensor : "
+                 "[-{},{}), got {}.",
+                 type(), output_nb_dims, output_nb_dims - 1, axis);
+    axes_rectified_idx.push_back(
+        static_cast<DimIdx_t>(axis >= 0 ? axis : axis + output_nb_dims));
+  }
+  // sort by descending order
+  std::sort(axes_rectified_idx.begin(), axes_rectified_idx.end());
+  // Raise error if duplicate indexes are found
+  const auto &it = std::adjacent_find(axes_rectified_idx.begin(), axes_rectified_idx.end());
+  AIDGE_ASSERT(
+      it == axes_rectified_idx.end(),
+      "{} : The index {} appears multiple times in list of input dims. "
+      "Check positive and negative indexes.\nRaw indexes :\t{}\nRectified "
+      "indexes :\t{}",
+      type(), *it, this->axes(), axes_rectified_idx);
+
+  // computation
+  std::vector<DimSize_t> output_dims(input_dims);
+  output_dims.reserve(input_dims.size() + this->axes().size());
+  for (const DimIdx_t &axis : axes_rectified_idx) {
+    output_dims.insert(output_dims.begin() + axis, 1);
+  }
+  mOutputs[0]->resize(output_dims);
+  return true;
+}
+
+void Unsqueeze_Op::setBackend(const std::string &name,
+                              Aidge::DeviceIdx_t device) {
+  if (Registrar<Unsqueeze_Op>::exists({name})) {
+    SET_IMPL_MACRO(Unsqueeze_Op, *this, name);
+  } else {
+    mImpl = std::make_shared<Unsqueeze_OpImpl>(*this);
+  }
+  mOutputs[0]->setBackend(name, device);
+}
+
+void Aidge::Unsqueeze_OpImpl::forward() {
+  const Unsqueeze_Op &op_ = static_cast<const Unsqueeze_Op &>(mOp);
+  // Check if input is provided
+  AIDGE_ASSERT(op_.getInput(0), "Unsqueeze : missing input 0");
+  op_.getOutput(0)->getImpl()->copy(op_.getInput(0)->getImpl()->rawPtr(),
+                                    op_.getInput(0)->size());
+}
+
+} // namespace Aidge
diff --git a/unit_tests/operator/Test_Unsqueeze_Op.cpp b/unit_tests/operator/Test_Unsqueeze_Op.cpp
new file mode 100644
index 000000000..79f5b89b1
--- /dev/null
+++ b/unit_tests/operator/Test_Unsqueeze_Op.cpp
@@ -0,0 +1,382 @@
+/********************************************************************************
+ * 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
+ *
+ ********************************************************************************/
+
+#include <algorithm>
+#include <chrono>
+#include <cmath>
+#include <cstddef> // std::size_t
+#include <cstdint> // std::uint16_t
+#include <fmt/core.h>
+#include <iostream>
+#include <memory>
+#include <numeric> // std::accumulate
+#include <ostream>
+#include <random> // std::random_device, std::mt19937, std::uniform_real_distribution
+#include <vector>
+
+#include <catch2/catch_test_macros.hpp>
+#include <catch2/generators/catch_generators_random.hpp>
+
+#include "aidge/data/Data.hpp"
+#include "aidge/data/Tensor.hpp"
+#include "aidge/operator/Unsqueeze.hpp"
+#include "aidge/utils/ArrayHelpers.hpp"
+#include "aidge/utils/TensorUtils.hpp"
+#include "aidge/utils/Types.h"
+
+namespace Aidge {
+
+bool ensure_axes_validity(std::vector<int8_t> dims_to_unsqueeze,
+                          DimIdx_t nb_dims_input_tensor) {
+
+  bool in_bounds =
+      std::all_of(dims_to_unsqueeze.begin(), dims_to_unsqueeze.end(),
+                  [&nb_dims_input_tensor,
+                   &dims_to_unsqueeze](const int8_t &dim_to_unsqueeze) {
+                    return (dim_to_unsqueeze <
+                            nb_dims_input_tensor + dims_to_unsqueeze.size());
+                  });
+
+  std::sort(dims_to_unsqueeze.begin(), dims_to_unsqueeze.end());
+  bool index_appear_twice =
+      dims_to_unsqueeze.end() !=
+      std::adjacent_find(dims_to_unsqueeze.begin(), dims_to_unsqueeze.end());
+
+  return in_bounds && !index_appear_twice;
+}
+
+std::vector<DimSize_t>
+generate_unsqueeze_output_dims(std::vector<size_t> dims_in,
+                               std::vector<int8_t> dims_to_unsqueeze) {
+
+  std::sort(dims_to_unsqueeze.begin(), dims_to_unsqueeze.end());
+  std::vector<DimSize_t> dims_out(dims_in);
+  dims_out.reserve(dims_in.size() + dims_to_unsqueeze.size());
+  for (const DimIdx_t &dim : dims_to_unsqueeze) {
+    dims_out.insert(dims_out.begin() + dim, 1);
+  }
+  return dims_out;
+}
+
+std::vector<int8_t> rectify_indexes(const std::vector<int8_t> & dims_to_unsqueeze,
+                                    const int8_t offset) {
+  std::vector<int8_t> output;
+  output.reserve(dims_to_unsqueeze.size());
+  for (int8_t dim : dims_to_unsqueeze) {
+    output.push_back(dim >= 0 ? dim : dim + offset);
+  }
+  return output;
+}
+
+TEST_CASE("[core/operator] Unsqueeze(forwardDims)",
+          "[Unsqueeze][forwardDims]") {
+  constexpr std::uint16_t NB_TRIALS = 10;
+  // Create a random number generator
+  auto random_seed = Catch::Generators::Detail::getSeed;
+  std::mt19937 gen(random_seed());
+
+  std::uniform_real_distribution<float> valueDist(0.1f, 1.1f);
+  std::uniform_int_distribution<std::size_t> tensor_dims_size_dist(
+      std::size_t(1), std::size_t(10));
+  std::uniform_int_distribution<std::size_t> tensor_nb_dims_dist(
+      std::size_t(1), std::size_t(7));
+  std::uniform_int_distribution<std::size_t> nb_dims_to_unsqueeze_dist(
+      std::size_t(1), std::size_t(8));
+
+  std::shared_ptr<Tensor> input_T = std::make_shared<Tensor>();
+  std::shared_ptr<Tensor> axes_T = std::make_shared<Tensor>();
+
+  SECTION("ERROR : Inputs not ready") {
+    SECTION("unconnected input") {
+      std::shared_ptr<Node> myUnsqueeze =
+          Unsqueeze(std::vector<std::int8_t>({0}));
+      auto op =
+          std::static_pointer_cast<OperatorTensor>(myUnsqueeze->getOperator());
+      REQUIRE_THROWS(op->forwardDims());
+    }
+
+    std::shared_ptr<Tensor> input_T = std::make_shared<Tensor>();
+
+    SECTION("empty tensor") {
+      // Create the Unsqueeze Operator
+      std::shared_ptr<Node> myUnsqueeze =
+          Unsqueeze(std::vector<std::int8_t>({0}));
+      auto op =
+          std::static_pointer_cast<OperatorTensor>(myUnsqueeze->getOperator());
+      op->associateInput(0, input_T);
+
+      CHECK(op->forwardDims() == false);
+    }
+  }
+  SECTION("Compare with reference output") {
+    int8_t nb_dims = 3;
+    SECTION("axes is given via attribute") {
+      SECTION("unsqueez(0)") {
+        std::shared_ptr<Node> myUnsqueeze =
+            Unsqueeze(std::vector<std::int8_t>({0}));
+        auto op = std::static_pointer_cast<OperatorTensor>(
+            myUnsqueeze->getOperator());
+        op->associateInput(0, input_T);
+
+        std::vector<DimSize_t> dims_in{2, 3, 4};
+        input_T->resize(dims_in);
+
+        CHECK(op->forwardDims() == true);
+        CHECK(op->getOutput(0)->dims() == std::vector<DimSize_t>({1, 2, 3, 4}));
+        CHECK((op->getOutput(0)->dims().size()) == nb_dims + 1);
+      }
+      SECTION("Unsqueeze(1)") {
+        std::shared_ptr<Node> myUnsqueeze =
+            Unsqueeze(std::vector<std::int8_t>({1}));
+        auto op = std::static_pointer_cast<OperatorTensor>(
+            myUnsqueeze->getOperator());
+        op->associateInput(0, input_T);
+
+        std::array<DimSize_t, 3> dims_in{2, 3, 4};
+        input_T->resize(dims_in);
+
+        CHECK(op->forwardDims() == true);
+        CHECK(op->getOutput(0)->dims() == std::vector<DimSize_t>({2, 1, 3, 4}));
+        CHECK((op->getOutput(0)->dims().size()) == nb_dims + 1);
+      }
+      SECTION("Unsqueeze(2)") {
+        std::shared_ptr<Node> myUnsqueeze =
+            Unsqueeze(std::vector<std::int8_t>({2}));
+        auto op = std::static_pointer_cast<OperatorTensor>(
+            myUnsqueeze->getOperator());
+        op->associateInput(0, input_T);
+
+        std::vector<DimSize_t> dims_in{2, 3, 4};
+        input_T->resize(dims_in);
+
+        CHECK(op->forwardDims() == true);
+        CHECK(op->getOutput(0)->dims() == std::vector<DimSize_t>({2, 3, 1, 4}));
+        CHECK((op->getOutput(0)->dims().size()) == nb_dims + 1);
+      }
+      SECTION("Unsqueeze({0,4})") {
+        std::shared_ptr<Node> myUnsqueeze =
+            Unsqueeze(std::vector<std::int8_t>({0, 4}));
+        auto op = std::static_pointer_cast<OperatorTensor>(
+            myUnsqueeze->getOperator());
+        op->associateInput(0, input_T);
+
+        std::vector<DimSize_t> dims_in{3, 4, 5};
+        input_T->resize(dims_in);
+
+        CHECK(op->forwardDims() == true);
+        CHECK(op->getOutput(0)->dims() ==
+              std::vector<DimSize_t>({1, 3, 4, 5, 1}));
+      }
+    }
+    SECTION("axes is given via tensor") {
+        // arguments here should be overriden by axes_T values
+        std::shared_ptr<Node> myUnsqueeze =
+            Unsqueeze(std::vector<std::int8_t>({0, 4}));
+        auto op = std::static_pointer_cast<OperatorTensor>(
+            myUnsqueeze->getOperator());
+        op->associateInput(0, input_T);
+
+        auto axes_T = std::make_shared<Aidge::Tensor>(
+            Aidge::Array1D<int8_t, 3>({1, 3, 4}));
+        axes_T->setDataType(Aidge::DataType::Int8);
+        axes_T->setBackend("cpu");
+
+        std::vector<DimSize_t> dims_in{3, 4, 5};
+        input_T->resize(dims_in);
+        op->associateInput(0, input_T);
+        op->associateInput(1, axes_T);
+
+        CHECK(op->forwardDims(true) == true);
+        CHECK(op->getOutput(0)->dims() ==
+              std::vector<DimSize_t>({3, 1, 4, 1, 1, 5}));
+    }
+  }
+  SECTION("Random testing") {
+    SECTION("Unsqueeze({N,...})") {
+      int number_of_operation{0};
+      for (uint16_t trial = 0; trial < NB_TRIALS; ++trial) {
+        const size_t nb_dims_to_unsqueeze = nb_dims_to_unsqueeze_dist(gen);
+        const size_t nb_dims_tensor = tensor_nb_dims_dist(gen);
+        const size_t idx_dims_to_unsqueeze_max =
+            nb_dims_to_unsqueeze + nb_dims_tensor;
+        const size_t variance_error = 2;
+        std::uniform_int_distribution<short> idx_dims_to_unsqueeze_dist(
+            -idx_dims_to_unsqueeze_max - variance_error,
+            idx_dims_to_unsqueeze_max - 1 + variance_error);
+        // Create the Operator
+        std::vector<int8_t> dims_to_unsqueeze(nb_dims_to_unsqueeze);
+        std::generate(dims_to_unsqueeze.begin(), dims_to_unsqueeze.end(),
+                      [&gen, &idx_dims_to_unsqueeze_dist]() {
+                        return idx_dims_to_unsqueeze_dist(gen);
+                      });
+        std::shared_ptr<Node> unsqueeze_node = Unsqueeze(dims_to_unsqueeze);
+        auto op = std::static_pointer_cast<OperatorTensor>(
+            unsqueeze_node->getOperator());
+        op->associateInput(0, input_T);
+
+        // input tensor
+        std::vector<std::size_t> dims_in(nb_dims_tensor);
+        std::generate(dims_in.begin(), dims_in.end(),
+                      [&gen, &tensor_dims_size_dist]() {
+                        return tensor_dims_size_dist(gen);
+                      });
+        input_T->resize(dims_in);
+        op->setInput(0, input_T);
+
+        dims_to_unsqueeze = rectify_indexes(
+            dims_to_unsqueeze, input_T->nbDims() + dims_to_unsqueeze.size());
+        bool dims_to_unsqueeze_valid =
+            ensure_axes_validity(dims_to_unsqueeze, input_T->nbDims());
+        Log::warn("raw dims_to_unsqueeze : {}", dims_to_unsqueeze);
+        Log::warn("dims_to_unsqueeze : {}", dims_to_unsqueeze);
+        Log::warn("tensor dims : {}", input_T->dims());
+
+        if (!dims_to_unsqueeze_valid) {
+          ensure_axes_validity(dims_to_unsqueeze, input_T->nbDims());
+          REQUIRE_THROWS(op->forwardDims(true));
+        } else {
+          // output tensor
+          std::vector<DimSize_t> dims_out =
+              generate_unsqueeze_output_dims(dims_in, dims_to_unsqueeze);
+          Log::warn("dims_out : {}", dims_out);
+          CHECK(op->forwardDims(true) == true);
+          CHECK(op->getOutput(0)->dims() == dims_out);
+          generate_unsqueeze_output_dims(dims_in, dims_to_unsqueeze);
+        }
+      }
+    }
+  }
+}
+
+TEST_CASE("[core/operator] Unsqueeze(forward)", "[Unsqueeze][forward]") {
+  constexpr std::uint16_t NB_TRIALS = 10;
+  // Create a random number generator
+  std::random_device rd;
+  auto random_seed = rd();
+  std::cout << "True random seed : " << random_seed << std::endl;
+  std::mt19937 gen(random_seed);
+  // Random float distribution between 0 and 1
+  std::uniform_real_distribution<float> valueDist(0.1f, 1.1f);
+  std::uniform_int_distribution<std::size_t> tensor_dims_size_dist(
+      std::size_t(1), std::size_t(10));
+  std::size_t min_tensor_nb_dims{1};
+  std::size_t max_tensor_nb_dims{7};
+  std::uniform_int_distribution<std::size_t> tensor_nb_dims_dist(
+      min_tensor_nb_dims, max_tensor_nb_dims);
+  std::uniform_int_distribution<std::size_t> nb_dims_to_unsqueeze_dist(
+      std::size_t(1), std::size_t(8));
+  std::uniform_int_distribution<short> idx_dims_to_unsqueeze_dist(-9, 8);
+
+  std::shared_ptr<Tensor> input_T = std::make_shared<Tensor>();
+  input_T->setDataType(DataType::Float32);
+  input_T->setBackend("cpu");
+  std::shared_ptr<Tensor> result_T = std::make_shared<Tensor>();
+  result_T->setDataType(DataType::Float32);
+  result_T->setBackend("cpu");
+
+  // BENCHMARKING
+  std::chrono::time_point<std::chrono::system_clock> start;
+  std::chrono::time_point<std::chrono::system_clock> end;
+  std::chrono::duration<double, std::micro> duration{};
+
+  int number_of_operation{0};
+  for (uint16_t trial = 0; trial < NB_TRIALS; ++trial) {
+    // Create the Operator
+    size_t nb_dims_to_unsqueeze = nb_dims_to_unsqueeze_dist(gen);
+    std::vector<int8_t> dims_to_unsqueeze(nb_dims_to_unsqueeze);
+    std::generate(dims_to_unsqueeze.begin(), dims_to_unsqueeze.end(),
+                  [&gen, &idx_dims_to_unsqueeze_dist]() {
+                    return idx_dims_to_unsqueeze_dist(gen);
+                  });
+    std::shared_ptr<Node> unsqueeze_node = Unsqueeze(dims_to_unsqueeze);
+    auto op =
+        std::static_pointer_cast<OperatorTensor>(unsqueeze_node->getOperator());
+    op->setDataType(DataType::Float32);
+    op->setBackend("cpu");
+    op->associateInput(0, input_T);
+
+    // input tensor
+    const std::size_t nb_dims_tensor = tensor_nb_dims_dist(gen);
+    std::vector<std::size_t> dims_in(nb_dims_tensor);
+    std::generate(dims_in.begin(), dims_in.end(),
+                  [&gen, &tensor_dims_size_dist]() {
+                    return tensor_dims_size_dist(gen);
+                  });
+    input_T->resize(dims_in);
+    op->setInput(0, input_T);
+
+    // rectifying indexes
+    std::transform(
+        dims_to_unsqueeze.begin(), dims_to_unsqueeze.end(),
+        dims_to_unsqueeze.begin(),
+        [&nb_dims_tensor, &nb_dims_to_unsqueeze](int8_t dim_to_unsqueeze) {
+          return dim_to_unsqueeze < 0
+                     ? dim_to_unsqueeze +
+                           (nb_dims_tensor + nb_dims_to_unsqueeze)
+                     : dim_to_unsqueeze;
+        });
+
+    // ensuring arguments given to Unsqueeze are good
+    bool axes_to_unsqueeze_valid =
+        ensure_axes_validity(dims_to_unsqueeze, input_T->nbDims());
+    if (!axes_to_unsqueeze_valid) {
+      REQUIRE_THROWS(op->forwardDims(true));
+    } else {
+      // output tensor
+      std::vector<DimSize_t> dims_out =
+          generate_unsqueeze_output_dims(dims_in, dims_to_unsqueeze);
+      CHECK(op->forwardDims(true) == true);
+      CHECK(op->getOutput(0)->dims() == dims_out);
+
+      SECTION("forward") {
+        const std::size_t nb_elems =
+            std::accumulate(dims_in.cbegin(), dims_in.cend(), std::size_t(1),
+                            std::multiplies<std::size_t>());
+        float *array_in = new float[nb_elems];
+        for (std::size_t i = 0; i < nb_elems; ++i) {
+          array_in[i] = valueDist(gen);
+        }
+        number_of_operation += nb_elems; // Copying all values : 1
+                                         // assignation / item in the tensor
+
+        // input0
+        input_T->resize(dims_in);
+        input_T->getImpl()->setRawPtr(array_in, nb_elems);
+
+        // results
+        result_T->resize(dims_out);
+        result_T->getImpl()->setRawPtr(array_in, nb_elems);
+
+        CHECK(op->forwardDims(true) == true);
+        start = std::chrono::system_clock::now();
+        REQUIRE_NOTHROW(unsqueeze_node->forward());
+        end = std::chrono::system_clock::now();
+        duration +=
+            std::chrono::duration_cast<std::chrono::microseconds>(end - start);
+
+        CHECK(result_T->nbDims() == op->getOutput(0)->nbDims());
+        for (DimSize_t i = 0; i < op->getOutput(0)->nbDims(); ++i) {
+          CHECK(result_T->dims().at(i) == op->getOutput(0)->dims().at(i));
+        }
+        CHECK(approxEq<float>(*result_T, *(op->getOutput(0))));
+
+        delete[] array_in;
+      }
+    }
+    std::cout << "Unsqueeze total execution time : " << duration.count() << "µs"
+              << std::endl;
+    std::cout << "Number of operations : " << number_of_operation << std::endl;
+    std::cout << "Operation / µs = " << number_of_operation / duration.count()
+              << std::endl;
+  }
+}
+
+} // namespace Aidge
-- 
GitLab