add dilations and ceilMode to AvgPooling

include/aidge/operator/AvgPooling.hpp

@@ -33,12 +33,21 @@ enum class AvgPoolingAttr {
      * Specifies the step size of the sliding window along each spatial dimension.
      */
     StrideDims,
-
+    /**
+     * @brief Dilation along each spatial axis. Default value is 1.
+     */
+    Dilations,
     /**
      * @brief Kernel dimensions for the pooling operation.
      * Specifies the size of the pooling window along each spatial dimension.
      */
-    KernelDims
+    KernelDims,
+    /**
+     * @brief Flag indicating whether to use ceil or floor when calculating output size.
+     * - `true`: Use `ceil` for output size calculation.
+     * - `false`: Use `floor` for output size calculation.
+     */
+    CeilMode
 };
 
 /**
@@ -67,7 +76,9 @@ private:
      */
     using Attributes_ = StaticAttributes<AvgPoolingAttr,
                                              std::array<DimSize_t, DIM>,
-                                             std::array<DimSize_t, DIM>>;
+                                             std::array<DimSize_t, DIM>,
+                                             std::array<DimSize_t, DIM>,
+                                             bool>;
     template <AvgPoolingAttr e>
     using attr = typename Attributes_::template attr<e>;
 
@@ -89,11 +100,15 @@ public:
      * Defaults to 1 for each dimension.
      */
     constexpr AvgPooling_Op(const std::array<DimSize_t, DIM> &kernel_dims,
-                            const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t, DIM>(1))
+                            const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t, DIM>(1),
+                            const std::array<DimSize_t, DIM> &dilations = create_array<DimSize_t, DIM>(1),
+                            bool ceil_mode = false)
         : OperatorTensor(Type, {InputCategory::Data}, 1),
           mAttributes(std::make_shared<Attributes_>(
                         attr<AvgPoolingAttr::StrideDims>(stride_dims),
-                        attr<AvgPoolingAttr::KernelDims>(kernel_dims)))
+                        attr<AvgPoolingAttr::KernelDims>(kernel_dims),
+                        attr<AvgPoolingAttr::Dilations>(dilations),
+                        attr<AvgPoolingAttr::CeilMode>(ceil_mode)))
     {}
 
     /**
@@ -155,11 +170,23 @@ public:
     inline std::array<DimSize_t, DIM>& strideDims() const { return mAttributes->template getAttr<AvgPoolingAttr::StrideDims>(); }
 
     /**
-     * @brief Accessor for the kernel dimensions.
-     * @return An array representing the kernel dimensions.
+     * @brief Accessor for dilations.
+     * @return An array representing spatial dilations.
+     */
+    inline std::array<DimSize_t, DIM>& dilations() const { return mAttributes->template getAttr<AvgPoolingAttr::Dilations>(); }
+
+    /**
+     * @brief Accessor for kernel dimensions.
+     * @return An array representing kernel dimensions.
      */
     inline std::array<DimSize_t, DIM>& kernelDims() const { return mAttributes->template getAttr<AvgPoolingAttr::KernelDims>(); }
 
+    /**
+     * @brief Accessor for ceil mode flag.
+     * @return Boolean value indicating whether ceil mode is enabled.
+     */
+    inline bool& ceilMode() const { return mAttributes->template getAttr<AvgPoolingAttr::CeilMode>(); }
+
     /**
      * @brief Retrieves the names of the input tensors.
      * @return A vector of strings representing the input tensors names.
@@ -180,31 +207,39 @@ public:
 /**
  * @brief Creates an AvgPooling operator node.
  * @tparam DIM Number of dimensions for the pooling operation.
- * @param kernel_dims Size of the pooling window for each spatial dimension.
- * @param name Name of the operator node. Defaults to an empty string.
- * @param stride_dims Step size (stride) for sliding the pooling window across the input dimensions. Defaults to 1 for each dimension.
+ * @param[in] kernel_dims Size of the pooling window for each spatial dimension.
+ * @param[in] name Name of the operator node. Defaults to an empty string.
+ * @param[in] stride_dims Step size (stride) for sliding the pooling window across the input dimensions. Defaults to 1 for each dimension.
+ * @param[in] dilations Spatial dilations for the pooling operation.
+ * @param[in] ceil_mode Indicates whether to use ceil mode for output size calculation.
  * @return A shared pointer to the created operator node.
  */
 template <std::array<DimSize_t, 1>::size_type DIM>
 std::shared_ptr<Node> AvgPooling(const std::array<DimSize_t, DIM> &kernel_dims,
                                  const std::string& name = "",
-                                 const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1));
+                                 const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+                                 const std::array<DimSize_t, DIM> &dilations = create_array<DimSize_t,DIM>(1),
+                                 bool ceil_mode=false);
 
 /**
  * @brief Overload of AvgPooling for C-style arrays.
  * @tparam DIM Number of dimensions for the pooling operation.
- * @param kernel_dims C-style array specifying the kernel dimensions.
- * @param name Name of the operator node. Defaults to an empty string.
- * @param stride_dims Step size (stride) for sliding the pooling window across the input dimensions. Defaults to 1 for each dimension.
+ * @param[in] kernel_dims C-style array specifying the kernel dimensions.
+ * @param[in] name Name of the operator node. Defaults to an empty string.
+ * @param[in] stride_dims Step size (stride) for sliding the pooling window across the input dimensions. Defaults to 1 for each dimension.
+ * @param[in] dilations Spatial dilations for the pooling operation.
+ * @param[in] ceil_mode Indicates whether to use ceil mode for output size calculation.
  * @return A shared pointer to the created operator node.
  */
 template <DimSize_t DIM>
 inline std::shared_ptr<Node> AvgPooling(
     DimSize_t const (&kernel_dims)[DIM],
     const std::string& name = "",
-    const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1)) {
+    const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+    const std::array<DimSize_t, DIM> &dilations = create_array<DimSize_t,DIM>(1),
+    bool ceil_mode=false) {
     static_assert(DIM<=MaxDim,"Too many kernel dimensions required by AvgPooling, not supported");
-    return AvgPooling(to_array(kernel_dims), name, stride_dims);
+    return AvgPooling(to_array(kernel_dims), name, stride_dims, dilations, ceil_mode);
 }
 }  // namespace Aidge
 
@@ -221,10 +256,7 @@ namespace {
  * @brief String representation of the AvgPooling attributes.
  */
 template <>
-const char *const EnumStrings<Aidge::AvgPoolingAttr>::data[] = {
-    "stride_dims",
-    "kernel_dims"
-};
+const char *const EnumStrings<Aidge::AvgPoolingAttr>::data[] = { "stride_dims", "kernel_dims", "dilations", "ceil_mode" };
 }
 
 #endif /* AIDGE_CORE_OPERATOR_AVGPOOLING_H_ */

include/aidge/operator/MetaOperatorDefs.hpp

@@ -164,6 +164,7 @@ PaddedConvDepthWise(const DimSize_t nb_channels,
  * @param[in] kernel_dims The dimensions of the pooling window.
  * @param[in] name Optional name for the operation.
  * @param[in] stride_dims The stride dimensions for pooling (default is 1).
+ * @param[in] dilations The spatial dilations for pooling (default is 1).
  * @param[in] padding_dims Padding dimensions before pooling (default is 0).
  * @return A shared pointer to the Node representing the padded average pooling operation.
  */
@@ -171,6 +172,7 @@ template <std::array<DimSize_t, 1>::size_type DIM>
 extern std::shared_ptr<Node> PaddedAvgPooling(const std::array<DimSize_t, DIM> &kernel_dims,
                                   const std::string& name = "",
                                   const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+                                  const std::array<DimSize_t, DIM> &dilations = create_array<DimSize_t,DIM>(1),
                                   const std::array<DimSize_t, 2*DIM> &padding_dims = create_array<DimSize_t,2*DIM>(0));
 
 /**
@@ -180,12 +182,14 @@ extern std::shared_ptr<Node> PaddedAvgPooling(const std::array<DimSize_t, DIM> &
  *
  * @param[in] kernel_dims The dimensions of the pooling window.
  * @param[in] stride_dims The stride dimensions for pooling (default is 1).
+ * @param[in] dilations The spatial dilations for pooling (default is 1).
  * @param[in] padding_dims Padding dimensions before pooling (default is 0).
  * @return A shared pointer to the MetaOperator_Op representing the padded average pooling operation.
  */
 template <std::array<DimSize_t, 1>::size_type DIM>
 extern std::shared_ptr<MetaOperator_Op> PaddedAvgPooling_Op(const std::array<DimSize_t, DIM> &kernel_dims,
                                   const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+                                  const std::array<DimSize_t, DIM> &dilations = create_array<DimSize_t,DIM>(1),
                                   const std::array<DimSize_t, 2*DIM> &padding_dims = create_array<DimSize_t,2*DIM>(0));
 
 // Helper function for average pooling with C-style array for kernel_dims, enabling automatic DIM deduction.
@@ -195,6 +199,8 @@ PaddedAvgPooling(DimSize_t const (&kernel_dims)[DIM],
                  const std::string &name = "",
                  const std::array<DimSize_t, DIM> &stride_dims =
                      create_array<DimSize_t, DIM>(1),
+                 const std::array<DimSize_t, DIM> &dilations =
+                     create_array<DimSize_t, DIM>(1),
                  const std::array<DimSize_t, 2 * DIM> &padding_dims =
                      create_array<DimSize_t, 2 * DIM>(0));
 
@@ -208,6 +214,7 @@ PaddedAvgPooling(DimSize_t const (&kernel_dims)[DIM],
  * @param[in] kernel_dims The dimensions of the pooling window.
  * @param[in] name Optional name for the operation.
  * @param[in] stride_dims The stride dimensions for pooling (default is 1).
+ * @param[in] dilations The spatial dilations for pooling (default is 1).
  * @param[in] padding_dims Padding dimensions before pooling (default is 0).
  * @param[in] ceil_mode Whether to use ceiling mode for pooling (default is false).
  * @return A shared pointer to the Node representing the padded max pooling operation.
@@ -216,11 +223,12 @@ template <std::array<DimSize_t, 1>::size_type DIM>
 inline std::shared_ptr<Node> PaddedMaxPooling(const std::array<DimSize_t, DIM> &kernel_dims,
                                   const std::string& name = "",
                                   const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+                                  const std::array<DimSize_t, DIM> &dilations = create_array<DimSize_t,DIM>(1),
                                   const std::array<DimSize_t, 2*DIM> &padding_dims = create_array<DimSize_t,2*DIM>(0),
                                   bool ceil_mode = false) {
     auto graph = Sequential({
         Pad<DIM>(padding_dims, (!name.empty()) ? name + "_pad" : ""),
-        MaxPooling(kernel_dims, (!name.empty()) ? name + "_maxpooling" : "", stride_dims, ceil_mode)
+        MaxPooling(kernel_dims, (!name.empty()) ? name + "_maxpooling" : "", stride_dims, dilations, ceil_mode)
     });
 
     return MetaOperator(("PaddedMaxPooling" + std::to_string(DIM) + "D").c_str(), graph, {}, name);
@@ -233,6 +241,7 @@ inline std::shared_ptr<Node> PaddedMaxPooling(const std::array<DimSize_t, DIM> &
  *
  * @param[in] kernel_dims The dimensions of the pooling window.
  * @param[in] stride_dims The stride dimensions for pooling (default is 1).
+ * @param[in] dilations The spatial dilations for pooling (default is 1).
  * @param[in] padding_dims Padding dimensions before pooling (default is 0).
  * @param[in] ceil_mode Whether to use ceiling mode for pooling (default is false).
  * @return A shared pointer to the MetaOperator_Op representing the padded max pooling operation.
@@ -240,11 +249,12 @@ inline std::shared_ptr<Node> PaddedMaxPooling(const std::array<DimSize_t, DIM> &
 template <std::array<DimSize_t, 1>::size_type DIM>
 inline std::shared_ptr<MetaOperator_Op> PaddedMaxPooling_Op(const std::array<DimSize_t, DIM> &kernel_dims,
                                   const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+                                  const std::array<DimSize_t, DIM> &dilations = create_array<DimSize_t,DIM>(1),
                                   const std::array<DimSize_t, 2*DIM> &padding_dims = create_array<DimSize_t,2*DIM>(0),
                                   bool ceil_mode = false) {
     auto graph = Sequential({
         Pad<DIM>(padding_dims, ""),
-        MaxPooling(kernel_dims, "", stride_dims, ceil_mode)
+        MaxPooling(kernel_dims, "", stride_dims, dilations, ceil_mode)
     });
     return std::make_shared<MetaOperator_Op>(("PaddedMaxPooling" + std::to_string(DIM) + "D").c_str(), graph);
 }
@@ -255,9 +265,10 @@ inline std::shared_ptr<Node> PaddedMaxPooling(
     DimSize_t const (&kernel_dims)[DIM],
     const std::string& name = "",
     const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+    const std::array<DimSize_t, DIM> &dilations = create_array<DimSize_t,DIM>(1),
     const std::array<DimSize_t, 2*DIM> &padding_dims = create_array<DimSize_t,2*DIM>(0),
     bool ceil_mode= false) {
-    return PaddedMaxPooling(to_array(kernel_dims), name, stride_dims, padding_dims, ceil_mode);
+    return PaddedMaxPooling(to_array(kernel_dims), name, stride_dims, padding_dims, dilations, ceil_mode);
 }
 
 /**

python_binding/operator/pybind_AvgPooling.cpp

@@ -47,11 +47,19 @@ template <DimIdx_t DIM> void declare_AvgPoolingOp(py::module &m) {
         :param stride_dims: The stride of the pooling operation. Specifies how much the kernel moves in each step.
                              By default, the stride is set to 1 for all dimensions.
         :type stride_dims: List[int], optional
+        :param dilations: The dilation value along each spatial axis of filter.
+        :type dilations: List[int], optional
+        :param ceil_mode: Whether to use ceil or floor when calculating the output dimensions.
+        :type ceil_mode: bool, optional
         )mydelimiter")
     .def(py::init<const std::array<DimSize_t, DIM> &,
-                  const std::array<DimSize_t, DIM> &>(),
+                  const std::array<DimSize_t, DIM> &,
+                  const std::array<DimSize_t, DIM> &,
+                  bool>(),
             py::arg("kernel_dims"),
-            py::arg("stride_dims") = create_array<DimSize_t, DIM>(1))
+            py::arg("stride_dims") = create_array<DimSize_t, DIM>(1),
+            py::arg("stride_dims") = create_array<DimSize_t, DIM>(1),
+            py::arg("ceil_mode") = false)
     .def("get_inputs_name", &AvgPooling_Op<DIM>::getInputsName)
     .def("get_outputs_name", &AvgPooling_Op<DIM>::getOutputsName)
     .def_readonly_static("Type", &AvgPooling_Op<DIM>::Type);
@@ -60,14 +68,19 @@ template <DimIdx_t DIM> void declare_AvgPoolingOp(py::module &m) {
 
   m.def(("AvgPooling" + std::to_string(DIM) + "D").c_str(), [](const std::vector<DimSize_t>& kernel_dims,
                                                                   const std::string& name,
-                                                                  const std::vector<DimSize_t>& stride_dims) {
+                                                                  const std::vector<DimSize_t>& stride_dims,
+                                                                  const std::vector<DimSize_t>& dilations,
+                                                                  bool ceil_mode) {
         AIDGE_ASSERT(kernel_dims.size() == DIM, "kernel_dims size [{}] does not match DIM [{}]", kernel_dims.size(), DIM);
         AIDGE_ASSERT(stride_dims.size() == DIM, "stride_dims size [{}] does not match DIM [{}]", stride_dims.size(), DIM);
+        AIDGE_ASSERT(dilations.size() == DIM, "dilations size [{}] does not match DIM [{}]", dilations.size(), DIM);
 
-        return AvgPooling<DIM>(to_array<DIM>(kernel_dims.begin()), name, to_array<DIM>(stride_dims.begin()));
+        return AvgPooling<DIM>(to_array<DIM>(kernel_dims.begin()), name, to_array<DIM>(stride_dims.begin()), to_array<DIM>(dilations.begin()), ceil_mode);
     }, py::arg("kernel_dims"),
        py::arg("name") = "",
        py::arg("stride_dims") = std::vector<DimSize_t>(DIM, 1),
+       py::arg("dilations") = std::vector<DimSize_t>(DIM, 1),
+       py::arg("ceil_mode") = false,
        R"mydelimiter(
         Initialize a node containing an AvgPooling operator.
 
@@ -75,6 +88,10 @@ template <DimIdx_t DIM> void declare_AvgPoolingOp(py::module &m) {
 
         :param kernel_dims: Size of the kernel applied during pooling.
         :type kernel_dims: List[int]
+        :param dilations: The dilation value along each spatial axis of filter.
+        :type dilations: List[int]
+        :param ceil_mode: Whether to use ceil or floor when calculating the output dimensions.
+        :type ceil_mode: bool
         :param name: Name of the operator node (optional).
         :type name: str
         :param stride_dims: Stride dimensions for the pooling operation.

src/operator/AvgPooling.cpp

@@ -47,17 +47,28 @@ std::shared_ptr<Aidge::Operator> Aidge::AvgPooling_Op<DIM>::clone() const {
 template <Aidge::DimIdx_t DIM>
 bool Aidge::AvgPooling_Op<DIM>::forwardDims(bool /*allowDataDependency*/) {
     if (inputsAssociated()) {
-        std::array<DimSize_t, DIM + 2> outputDims;
+        std::array<DimSize_t, DIM + 2> outputDims{};
         const std::array<DimSize_t, DIM + 2> inputDims(getInput(0)->template dims<DIM+2>());
-        outputDims[0] = inputDims[0];
-        outputDims[1] = inputDims[1];
 
-        for (std::size_t dim = 0; dim < mAttributes->template getAttr<AvgPoolingAttr::KernelDims>().size() ; ++dim) {
+        std::function<float(float)> roundingFunction;
+        if (mAttributes->template getAttr<AvgPoolingAttr::CeilMode>()) {
+            roundingFunction = [](float x) { return std::ceil(x); };
+        } else {
+            roundingFunction = [](float x) { return std::floor(x); };
+        }
+
+        for (std::size_t dim = 0; dim < mAttributes->template getAttr<AvgPoolingAttr::KernelDims>().size(); ++dim) {
+            const auto kernelDim = mAttributes->template getAttr<AvgPoolingAttr::KernelDims>()[dim];
+            const auto strideDim = mAttributes->template getAttr<AvgPoolingAttr::StrideDims>()[dim];
+            const auto dilationDim = mAttributes->template getAttr<AvgPoolingAttr::Dilations>()[dim];
+
             outputDims[dim+2] = 1 + static_cast<DimSize_t>(
-                                        std::floor(static_cast<float>(inputDims[dim+2] -
-                                                            mAttributes->template getAttr<AvgPoolingAttr::KernelDims>()[dim]) /
-                                        static_cast<float>(mAttributes->template getAttr<AvgPoolingAttr::StrideDims>()[dim])));
+                                            roundingFunction(static_cast<float>(inputDims[dim+2] -
+                                                                    (kernelDim - 1) * dilationDim - 1) /
+                                            static_cast<float>(strideDim)));
         }
+        outputDims[1] = inputDims[1];
+        outputDims[0] = inputDims[0];
         getOutput(0)->resize(outputDims);
         return true;
     }
@@ -128,10 +139,12 @@ template class Aidge::AvgPooling_Op<4>;
 template <std::array<Aidge::DimSize_t, 1>::size_type DIM>
 std::shared_ptr<Aidge::Node> Aidge::AvgPooling(const std::array<Aidge::DimSize_t, DIM> &kernel_dims,
                                            const std::string& name,
-                                           const std::array<Aidge::DimSize_t, DIM> &stride_dims) {
+                                           const std::array<Aidge::DimSize_t, DIM> &stride_dims,
+                                           const std::array<Aidge::DimSize_t, DIM> &dilations,
+                                           bool ceil_mode) {
     AIDGE_ASSERT(DIM<=MaxDim, "Too many kernel dimensions required by {}, not supported", AvgPooling_Op<DIM>::Type);
-    return std::make_shared<Node>(std::make_shared<AvgPooling_Op<static_cast<DimIdx_t>(DIM)>>(kernel_dims, stride_dims), name);
+    return std::make_shared<Node>(std::make_shared<AvgPooling_Op<static_cast<DimIdx_t>(DIM)>>(kernel_dims, stride_dims, dilations, ceil_mode), name);
 }
-template std::shared_ptr<Aidge::Node> Aidge::AvgPooling<1>(const std::array<Aidge::DimSize_t, 1>&, const std::string&, const std::array<Aidge::DimSize_t, 1>&);
-template std::shared_ptr<Aidge::Node> Aidge::AvgPooling<2>(const std::array<Aidge::DimSize_t, 2>&, const std::string&, const std::array<Aidge::DimSize_t, 2>&);
-template std::shared_ptr<Aidge::Node> Aidge::AvgPooling<3>(const std::array<Aidge::DimSize_t, 3>&, const std::string&, const std::array<Aidge::DimSize_t, 3>&);
+template std::shared_ptr<Aidge::Node> Aidge::AvgPooling<1>(const std::array<Aidge::DimSize_t, 1>&, const std::string&, const std::array<Aidge::DimSize_t, 1>&, const std::array<Aidge::DimSize_t, 1>&, bool);
+template std::shared_ptr<Aidge::Node> Aidge::AvgPooling<2>(const std::array<Aidge::DimSize_t, 2>&, const std::string&, const std::array<Aidge::DimSize_t, 2>&, const std::array<Aidge::DimSize_t, 2>&, bool);
+template std::shared_ptr<Aidge::Node> Aidge::AvgPooling<3>(const std::array<Aidge::DimSize_t, 3>&, const std::string&, const std::array<Aidge::DimSize_t, 3>&, const std::array<Aidge::DimSize_t, 3>&, bool);