Add a Stack operator

Context

The Stack operator is designed to be the counterpart of the Pop operator, which dictates its functionalities: at each time step t, it takes a tensor as input and stacks it along axis 0. The Pop operator overrides the getNbRequiredData() method, which specifies the amount of data consumed each time the operator is executed. Similarly, the Stack operator overrides the getNbRequiredMemory() method, which specifies the amount of data produced by the operator with each execution.

A notable difference from the Pop operator is the way we define the output tensor's dimensions (which has to be done statically). For Pop, this is straightforward: we simply take the input tensor's dimensions and remove the first dimension. For Stack, however, we define an attribute representing the maximum number of elements to stack. Then, given the input tensor shape, the output tensor will have an additional first dimension corresponding to the maximum number of elements to stack.

Of course, care must be taken during the scheduling phase and execution to ensure that the operator is not invoked more times than the defined maximum.

Modified files

• Stack.hpp and Stack.cpp : Cpp implementation;
• Test_StackImpl.cpp : Cpp unit tests;
• pybind_core.cpp and pybind_Stack.cpp : Added python bindings;

Detailed major modifications

TODO

• C++ Implementation
• Tests
• Python bindings

Scheduling tests are needed, but will be done in aidge_backend_cpu, as to have more flexibility (in particular, it would be interesting to test something like Pop->Dummy MetaOp->Stack and see how it goes).

src/operator/Stack.cpp

@@ -22,6 +22,8 @@
 
 namespace Aidge {
 
+// TODO: Check why getRequiredMemory is always called with empty vector as
+// inputSize
 Aidge::Elts_t Aidge::StackProdConso::getRequiredMemory(
     const Aidge::IOIndex_t inputIdx,
     const std::vector<DimSize_t> &inputsSize) const {
@@ -39,7 +41,8 @@ void StackOpImpl::forward() {
     const StackOp &op = dynamic_cast<const StackOp &>(mOp);
     AIDGE_ASSERT(op.getInput(0), "missing input #0");
     AIDGE_ASSERT((op.forwardStep() < op.maxElements()),
-                 "cannot forward anymore, maximum number of elements to stack exceeded");
+                 "cannot forward anymore, maximum number of elements to stack "
+                 "exceeded");
 
     op.getOutput(0)->getImpl()->copy(
         op.getInput(0)->getImpl()->rawPtr(),
@@ -53,7 +56,9 @@ StackOp::StackOp(std::uint32_t maxElements)
           attr<StackAttr::MaxElements>(maxElements),
           attr<StackAttr::ForwardStep>(0))) {
     if (maxElements == 0) {
-        AIDGE_THROW_OR_ABORT(std::invalid_argument, "StackOp creation failed: maxElements must be greater than 0.");
+        AIDGE_THROW_OR_ABORT(
+            std::invalid_argument,
+            "StackOp creation failed: maxElements must be greater than 0.");
     }
     mImpl = std::make_shared<StackOpImpl>(*this);
 }