src/armnnTfLiteParser/test/Transpose.cpp

//
// Copyright © 2019 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//

#include <boost/test/unit_test.hpp>
#include "ParserFlatbuffersFixture.hpp"
#include "../TfLiteParser.hpp"

BOOST_AUTO_TEST_SUITE(TensorflowLiteParser)

struct TransposeFixture : public ParserFlatbuffersFixture
{
    explicit TransposeFixture(const std::string & inputShape,
                              const std::string & permuteData,
                              const std::string & outputShape)
    {
        m_JsonString = R"(
            {
                  "version": 3,
                  "operator_codes": [
                    {
                      "builtin_code": "TRANSPOSE",
                      "version": 1
                    }
                  ],
                  "subgraphs": [
                    {
                      "tensors": [
                        {
                          "shape": )" + inputShape + R"(,
                          "type": "FLOAT32",
                          "buffer": 0,
                          "name": "inputTensor",
                          "quantization": {
                            "min": [
                              0.0
                            ],
                            "max": [
                              255.0
                            ],
                            "details_type": 0,
                            "quantized_dimension": 0
                          },
                          "is_variable": false
                        },
                        {
                          "shape": )" + outputShape + R"(,
                          "type": "FLOAT32",
                          "buffer": 1,
                          "name": "outputTensor",
                          "quantization": {
                            "details_type": 0,
                            "quantized_dimension": 0
                          },
                          "is_variable": false
                        })";
        if (!permuteData.empty())
        {
            m_JsonString += R"(,
                              {
                                "shape": [
                                  3
                                ],
                                "type": "INT32",
                                "buffer": 2,
                                "name": "permuteTensor",
                                "quantization": {
                                  "details_type": 0,
                                  "quantized_dimension": 0
                                },
                                "is_variable": false
                              })";
        }

        m_JsonString += R"(],
                      "inputs": [
                        0
                      ],
                      "outputs": [
                        1
                      ],
                      "operators": [
                        {
                          "opcode_index": 0,
                          "inputs": [
                            0)";
        if (!permuteData.empty())
        {
            m_JsonString += R"(,2)";
        }
        m_JsonString += R"(],
                          "outputs": [
                            1
                          ],
                          "builtin_options_type": "TransposeOptions",
                          "builtin_options": {
                          },
                          "custom_options_format": "FLEXBUFFERS"
                        }
                      ]
                    }
                  ],
                  "description": "TOCO Converted.",
                  "buffers": [
                    { },
                    { })";
        if (!permuteData.empty())
        {
            m_JsonString += R"(,{"data": )" + permuteData + R"( })";
        }
        m_JsonString += R"(
                  ]
                }
        )";
        Setup();
    }
};

struct TransposeFixtureWithPermuteData : TransposeFixture
{
    TransposeFixtureWithPermuteData() : TransposeFixture("[ 2, 2, 3 ]",
                                                         "[ 0, 0, 0, 0, 2, 0, 0, 0, 1, 0, 0, 0 ]",
                                                         "[ 2, 3, 2 ]") {}
};

BOOST_FIXTURE_TEST_CASE(TransposeWithPermuteData, TransposeFixtureWithPermuteData)
{
    RunTest<3, armnn::DataType::Float32>(
      0,
      {{"inputTensor", {  1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 }}},
      {{"outputTensor", {  1, 4, 2, 5, 3, 6, 7, 10, 8, 11, 9, 12 }}});

    BOOST_TEST((m_Parser->GetNetworkOutputBindingInfo(0, "outputTensor").second.GetShape()
                == armnn::TensorShape({2,3,2})));
}

struct TransposeFixtureWithoutPermuteData : TransposeFixture
{
    TransposeFixtureWithoutPermuteData() : TransposeFixture("[ 2, 2, 3 ]",
                                                            "",
                                                            "[ 2, 3, 2 ]") {}
};

BOOST_FIXTURE_TEST_CASE(TransposeWithoutPermuteDims, TransposeFixtureWithoutPermuteData)
{
    RunTest<3, armnn::DataType::Float32>(
        0,
        {{"inputTensor", {  1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 }}},
        {{"outputTensor", {  1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 }}});

    BOOST_TEST((m_Parser->GetNetworkOutputBindingInfo(0, "outputTensor").second.GetShape()
                == armnn::TensorShape({2,3,2})));
}

BOOST_AUTO_TEST_SUITE_END()