Test_ExpandImpl.cpp

/********************************************************************************
 * Copyright (c) 2024 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 <memory>

#include <catch2/catch_test_macros.hpp>

#include "aidge/data/DataType.hpp"
#include "aidge/data/Tensor.hpp"
#include "aidge/operator/Expand.hpp"
#include "aidge/utils/ArrayHelpers.hpp"


namespace Aidge {

using std::shared_ptr;

static void setupTestExpand(shared_ptr<Tensor> inputData,
                            shared_ptr<Tensor> inputShape,
                            shared_ptr<Expand_Op> &op,
                            Tensor &expectedOutput) {

    op->getOutput(0)->setDataType(inputData->dataType());

    inputData->setBackend("cpu");
    op->associateInput(0, inputData);

    inputShape->setBackend("cpu");
    op->associateInput(1, inputShape);

    expectedOutput.setBackend("cpu");
    expectedOutput.setDataType(DataType::Int32);
}

TEST_CASE("[cpu/operator] Expand(forward)", "[Expand][CPU]") {
    std::shared_ptr<Expand_Op> op = std::make_shared<Expand_Op>();
    op->setBackend("cpu");

    SECTION("Expand shape is bigger than inputData") {
        auto inputData = std::make_shared<Tensor>(Array1D<int, 2>({1, 3}));
        auto inputShape =
            std::make_shared<Tensor>(Array1D<std::int64_t, 4>({1, 3, 4, 2}));
        Tensor expectedOutput =
            Array4D<cpptype_t<DataType::Int32>, 1, 3, 4, 2>({{{{{1, 3}, {1, 3}, {1, 3}, {1, 3}},
                                        {{1, 3}, {1, 3}, {1, 3}, {1, 3}},
                                        {{1, 3}, {1, 3}, {1, 3}, {1, 3}}}}});
        setupTestExpand(inputData, inputShape, op, expectedOutput);

        // forwardDims has already been tested in core
        CHECK(op->forwardDims(true));
        REQUIRE_NOTHROW(op->forward());
        REQUIRE(expectedOutput == *op->getOutput(0));
    }
    SECTION("Expand shape has less dimensions than inputData") {
        auto inputData = std::make_shared<Tensor>(
            Array3D<int, 2, 1, 3>({{{2, 1, 3}, {2, 1, 3}}}));
        auto inputShape =
            std::make_shared<Tensor>(Array1D<std::int64_t, 2>({2, 3}));
        Tensor expectedOutput = Array3D<cpptype_t<DataType::Int32>, 2, 2, 3>(
            {{{{2, 1, 3}, {2, 1, 3}}, {{2, 1, 3}, {2, 1, 3}}}});
        setupTestExpand(inputData, inputShape, op,expectedOutput);

        // forwardDims has already been tested in core
        CHECK(op->forwardDims(true));
        REQUIRE_NOTHROW(op->forward());
        REQUIRE(expectedOutput == *op->getOutput(0));
    }
    SECTION("Expand shape = {1} leads to input equal to output.") {
        auto inputData = std::make_shared<Tensor>(
            Array4D<int, 2, 1, 3, 1>({{{2, 1, 3}, {2, 1, 3}}}));
        auto inputShape =
            std::make_shared<Tensor>(Array1D<std::int64_t, 1>({1}));
        Tensor expectedOutput =
            Array4D<cpptype_t<DataType::Int32>, 2, 1, 3, 1>({{{2, 1, 3}, {2, 1, 3}}});
        setupTestExpand(inputData, inputShape, op, expectedOutput);

        // forwardDims has already been tested in core
        CHECK(op->forwardDims(true));
        REQUIRE_NOTHROW(op->forward());
        REQUIRE(expectedOutput == *op->getOutput(0));
    }
    SECTION("The only common dimension is the last one & its equal to 1") {
        auto inputData = std::make_shared<Tensor>(
            Array4D<int, 1, 1, 3, 1>({{{{2, 1, 3}}}}));
        auto inputShape =
            std::make_shared<Tensor>(Array1D<std::int64_t, 3>({2, 1, 1}));
        Tensor expectedOutput =
            Array4D<cpptype_t<DataType::Int32>, 1, 2, 3, 1>({{{{2, 1, 3}, {2, 1, 3}}}});
        setupTestExpand(inputData, inputShape, op,expectedOutput);

        // forwardDims has already been tested in core
        CHECK(op->forwardDims(true));
        REQUIRE_NOTHROW(op->forward());
        REQUIRE(expectedOutput == *op->getOutput(0));
    }
    SECTION("N-Dim to N-Dim") {}
    auto inputData = std::shared_ptr<Tensor>();
}
} // namespace Aidge