diff --git a/include/aidge/operator/ReduceMean.hpp b/include/aidge/operator/ReduceMean.hpp index df1b46482558cc803455ce33ee18d8ec67f413da..8d3ed2f7117e7f375471f7c3f0cbe95555640a84 100644 --- a/include/aidge/operator/ReduceMean.hpp +++ b/include/aidge/operator/ReduceMean.hpp @@ -31,18 +31,18 @@ enum class ReduceMeanAttr { Axes, KeepDims }; template <DimIdx_t DIM> class ReduceMean_Op : public OperatorTensor, public Registrable<ReduceMean_Op<DIM>, std::string, std::unique_ptr<OperatorImpl>(const ReduceMean_Op<DIM> &)>, - public StaticAttributes<ReduceMeanAttr, std::array<DimSize_t, DIM>, DimSize_t> { + public StaticAttributes<ReduceMeanAttr, std::array<int, DIM>, DimSize_t> { public: static constexpr const char *Type = "ReduceMean"; ReduceMean_Op() = delete; - using Attributes_ = StaticAttributes<ReduceMeanAttr, std::array<DimSize_t, DIM>, DimSize_t>; + using Attributes_ = StaticAttributes<ReduceMeanAttr, std::array<int, DIM>, DimSize_t>; template <ReduceMeanAttr e> using attr = typename Attributes_::template attr<e>; - constexpr ReduceMean_Op(const std::array<DimSize_t, DIM> &axes, DimSize_t keep_dims) + constexpr ReduceMean_Op(const std::array<int, DIM> &axes, DimSize_t keep_dims) : OperatorTensor(Type, 1, 0, 1), Attributes_(attr<ReduceMeanAttr::Axes>(axes), attr<ReduceMeanAttr::KeepDims>(keep_dims)) {} @@ -74,7 +74,9 @@ class ReduceMean_Op : public OperatorTensor, bool reducedDim = false; for(std::size_t i=0; i<DIM; ++i) { - if(this->template getAttr<ReduceMeanAttr::Axes>()[i] == d) + int axis_ = this->template getAttr<ReduceMeanAttr::Axes>()[i]; + std::size_t axis= axis_>=0? axis_: axis_ + getInput(0)->nbDims(); + if(axis == d) { reducedDim = true; break; @@ -87,8 +89,11 @@ class ReduceMean_Op : public OperatorTensor, } else outDims.push_back(getInput(0)->dims()[d]); - } - mOutputs[0]->resize(outDims); + } + if(outDims.size()>0) + mOutputs[0]->resize(outDims); + else + mOutputs[0]->resize({1}); } } @@ -109,8 +114,8 @@ class ReduceMean_Op : public OperatorTensor, }; template <std::array<DimSize_t, 1>::size_type DIM> -inline std::shared_ptr<Node> ReduceMean(const std::array<DimSize_t, DIM> &axes, - DimSize_t keep_dims, +inline std::shared_ptr<Node> ReduceMean(const std::array<int, DIM> &axes, + DimSize_t keep_dims=1, const std::string& name = "") { // FIXME: properly handle default w&b initialization in every cases static_assert(DIM<=MaxDim,"Too many kernel dimensions required by ReduceMean, not supported"); @@ -121,7 +126,7 @@ inline std::shared_ptr<Node> ReduceMean(const std::array<DimSize_t, DIM> &axes, // helper with C-style array instead of std::array for kernel_dims to allow automatic template DIM deduction template <DimSize_t DIM> inline std::shared_ptr<Node> ReduceMean( - DimSize_t const (&axes)[DIM], + int const (&axes)[DIM], DimSize_t keep_dims = 1, const std::string& name = "") { static_assert(DIM<=MaxDim,"Too many kernel dimensions required by ReduceMean, not supported"); diff --git a/include/aidge/operator/Reshape.hpp b/include/aidge/operator/Reshape.hpp index 2d9372c4e5a23fab72df923e4408714255cee421..368f4fff3f85aacbf68f3225b0a83410c35b14f8 100644 --- a/include/aidge/operator/Reshape.hpp +++ b/include/aidge/operator/Reshape.hpp @@ -32,7 +32,7 @@ class Reshape_Op : public OperatorTensor, public: static constexpr const char* Type = "Reshape"; - Reshape_Op() : OperatorTensor(Type, 2, 0, 1) {} //1,1,1 + Reshape_Op() : OperatorTensor(Type, 2, 0, 1) {} /** * @brief Copy-constructor. Copy the operator attributes and its output tensor(s), but not its input tensors (the new operator has no input associated). diff --git a/python_binding/operator/pybind_ReduceMean.cpp b/python_binding/operator/pybind_ReduceMean.cpp index 1d0bcf3239f760ab1867c898d3d131b1dfe318a3..e5de98b69adde5133dde302f7306bc8a5c471eef 100644 --- a/python_binding/operator/pybind_ReduceMean.cpp +++ b/python_binding/operator/pybind_ReduceMean.cpp @@ -30,7 +30,7 @@ template <DimIdx_t DIM> void declare_ReduceMeanOp(py::module &m) { .def("get_outputs_name", &ReduceMean_Op<DIM>::getOutputsName) ; - m.def(("ReduceMean" + std::to_string(DIM) + "D").c_str(), [](const std::vector<DimSize_t>& axes, + m.def(("ReduceMean" + std::to_string(DIM) + "D").c_str(), [](const std::vector<int>& axes, DimSize_t keepDims, const std::string& name) { AIDGE_ASSERT(axes.size() == DIM, "axes size [%ld] does not match DIM [%d]", axes.size(), DIM);