ReduceMean.cpp

/********************************************************************************
 * 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/ReduceMean.hpp"

#include <algorithm>  // std::for_each, std::sort
#include <cstddef>    // std::size_t
#include <cstdint>    // std::int32_t
#include <memory>
#include <numeric> // For std::iota
#include <stdexcept>  // std::runtime_error
#include <string>
#include <vector>

#include "aidge/data/Tensor.hpp"
#include "aidge/utils/ErrorHandling.hpp"
#include "aidge/utils/Registrar.hpp"
#include "aidge/utils/Types.h"

const std::string Aidge::ReduceMean_Op::Type = "ReduceMean";

Aidge::ReduceMean_Op::ReduceMean_Op(const std::vector<std::int32_t>& axes, bool keep_dims, bool noop_with_empty_axes)
    : OperatorTensor(Type, {InputCategory::Data}, 1),
        mAttributes(std::make_shared<Attributes_>(
        attr<ReduceMeanAttr::Axes>(axes),
        attr<ReduceMeanAttr::KeepDims>(keep_dims),
        attr<ReduceMeanAttr::NoopWithEmptyAxes>(noop_with_empty_axes)))
{}

Aidge::ReduceMean_Op::ReduceMean_Op(const Aidge::ReduceMean_Op& op)
    : OperatorTensor(op),
        mAttributes(op.mAttributes)
{
    if (op.mImpl){
        SET_IMPL_MACRO(ReduceMean_Op, *this, op.backend());
    } else {
        mImpl = nullptr;
    }
}

std::shared_ptr<Aidge::Operator> Aidge::ReduceMean_Op::clone() const {
    return std::make_shared<ReduceMean_Op>(*this);
}

bool Aidge::ReduceMean_Op::forwardDims(bool /*allowDataDependency*/) {
    if (inputsAssociated()) {
        // make Axes attribute positive
        std::vector<std::int32_t>& axes = mAttributes->template getAttr<ReduceMeanAttr::Axes>();
        std::for_each(axes.begin(), axes.end(), [&] (std::int32_t& val) {
            if (val < 0)
                val+=static_cast<std::int32_t>(getInput(0)->nbDims());
        });
        std::sort(axes.begin(), axes.end());

        // build output dimensions
        std::vector<DimSize_t> outDims = getInput(0)->dims();

        if (axes.empty())
        {
            if(mAttributes->template getAttr<ReduceMeanAttr::NoopWithEmptyAxes>()) {
                mOutputs[0]->resize(outDims);
                return true;
            }
            // if no axes are provided and NoopWithEmptyAxes is false, reduce on all axes
            axes.resize(getInput(0)->nbDims());
            std::iota(axes.begin(), axes.end(), 0);
        }

        if (mAttributes->template getAttr<ReduceMeanAttr::KeepDims>()) {
            std::for_each(axes.cbegin(), axes.cend(), [&outDims] (const std::int32_t& val) { outDims[val] = 1; });
        }
        else {
            for (auto it = axes.crbegin(); it != axes.crend(); ++it)
                outDims.erase(outDims.begin() + static_cast<std::size_t>(*it));
        }

        // TODO: change {1} for {} when scalar Tensors are better handled.
        mOutputs[0]->resize((outDims.size()>0) ? outDims : std::vector<DimSize_t>({1}));
        return true;
    }
    return false;
}

void Aidge::ReduceMean_Op::setBackend(const std::string& name, Aidge::DeviceIdx_t device) {
    SET_IMPL_MACRO(ReduceMean_Op, *this, name);
    mOutputs[0]->setBackend(name, device);
}

std::set<std::string> Aidge::ReduceMean_Op::getAvailableBackends() const {
    return Registrar<ReduceMean_Op>::getKeys();
}

Aidge::ReduceMean_Op::~ReduceMean_Op() noexcept = default;

////////////////////////////////////////////

std::shared_ptr<Aidge::Node> Aidge::ReduceMean(const std::vector<std::int32_t> &axes,
                                        bool keep_dims,
                                        bool noop_with_empty_axes,
                                        const std::string& name) {
    AIDGE_ASSERT(axes.size()<=MaxDim, "Too many kernel dimensions required by ReduceMean, not supported");
    return std::make_shared<Node>(std::make_shared<ReduceMean_Op>(axes, keep_dims, noop_with_empty_axes), name);
}