text: //
  // Copyright © 2019 Arm Ltd. All rights reserved.
  // SPDX-License-Identifier: MIT
  //
  
  #include <armnn/Types.hpp>
  
  #include <armnnUtils/TensorUtils.hpp>
  
  #include <boost/test/unit_test.hpp>
  
  using namespace armnn;
  using namespace armnnUtils;
  
  BOOST_AUTO_TEST_SUITE(TensorUtilsSuite)
  
  BOOST_AUTO_TEST_CASE(ExpandDimsAxis0Test)
  {
      armnn::TensorShape inputShape({ 2, 3, 4 });
  
      // Expand dimension 0
      armnn::TensorShape outputShape = ExpandDims(inputShape, 0);
      BOOST_TEST(outputShape.GetNumDimensions() == 4);
      BOOST_TEST(outputShape[0] == 1);
      BOOST_TEST(outputShape[1] == 2);
      BOOST_TEST(outputShape[2] == 3);
      BOOST_TEST(outputShape[3] == 4);
  }
  
  BOOST_AUTO_TEST_CASE(ExpandDimsAxis1Test)
  {
      armnn::TensorShape inputShape({ 2, 3, 4 });
  
      // Expand dimension 1
      armnn::TensorShape outputShape = ExpandDims(inputShape, 1);
      BOOST_TEST(outputShape.GetNumDimensions() == 4);
      BOOST_TEST(outputShape[0] == 2);
      BOOST_TEST(outputShape[1] == 1);
      BOOST_TEST(outputShape[2] == 3);
      BOOST_TEST(outputShape[3] == 4);
  }
  
  BOOST_AUTO_TEST_CASE(ExpandDimsAxis2Test)
  {
      armnn::TensorShape inputShape({ 2, 3, 4 });
  
      // Expand dimension 2
      armnn::TensorShape outputShape = ExpandDims(inputShape, 2);
      BOOST_TEST(outputShape.GetNumDimensions() == 4);
      BOOST_TEST(outputShape[0] == 2);
      BOOST_TEST(outputShape[1] == 3);
      BOOST_TEST(outputShape[2] == 1);
      BOOST_TEST(outputShape[3] == 4);
  }
  
  BOOST_AUTO_TEST_CASE(ExpandDimsAxis3Test)
  {
      armnn::TensorShape inputShape({ 2, 3, 4 });
  
      // Expand dimension 3
      armnn::TensorShape outputShape = ExpandDims(inputShape, 3);
      BOOST_TEST(outputShape.GetNumDimensions() == 4);
      BOOST_TEST(outputShape[0] == 2);
      BOOST_TEST(outputShape[1] == 3);
      BOOST_TEST(outputShape[2] == 4);
      BOOST_TEST(outputShape[3] == 1);
  }
  
  BOOST_AUTO_TEST_CASE(ExpandDimsNegativeAxis1Test)
  {
      armnn::TensorShape inputShape({ 2, 3, 4 });
  
      // Expand dimension -1
      armnn::TensorShape outputShape = ExpandDims(inputShape, -1);
      BOOST_TEST(outputShape.GetNumDimensions() == 4);
      BOOST_TEST(outputShape[0] == 2);
      BOOST_TEST(outputShape[1] == 3);
      BOOST_TEST(outputShape[2] == 4);
      BOOST_TEST(outputShape[3] == 1);
  }
  
  BOOST_AUTO_TEST_CASE(ExpandDimsNegativeAxis2Test)
  {
      armnn::TensorShape inputShape({ 2, 3, 4 });
  
      // Expand dimension -2
      armnn::TensorShape outputShape = ExpandDims(inputShape, -2);
      BOOST_TEST(outputShape.GetNumDimensions() == 4);
      BOOST_TEST(outputShape[0] == 2);
      BOOST_TEST(outputShape[1] == 3);
      BOOST_TEST(outputShape[2] == 1);
      BOOST_TEST(outputShape[3] == 4);
  }
  
  BOOST_AUTO_TEST_CASE(ExpandDimsNegativeAxis3Test)
  {
      armnn::TensorShape inputShape({ 2, 3, 4 });
  
      // Expand dimension -3
      armnn::TensorShape outputShape = ExpandDims(inputShape, -3);
      BOOST_TEST(outputShape.GetNumDimensions() == 4);
      BOOST_TEST(outputShape[0] == 2);
      BOOST_TEST(outputShape[1] == 1);
      BOOST_TEST(outputShape[2] == 3);
      BOOST_TEST(outputShape[3] == 4);
  }
  
  BOOST_AUTO_TEST_CASE(ExpandDimsNegativeAxis4Test)
  {
      armnn::TensorShape inputShape({ 2, 3, 4 });
  
      // Expand dimension -4
      armnn::TensorShape outputShape = ExpandDims(inputShape, -4);
      BOOST_TEST(outputShape.GetNumDimensions() == 4);
      BOOST_TEST(outputShape[0] == 1);
      BOOST_TEST(outputShape[1] == 2);
      BOOST_TEST(outputShape[2] == 3);
      BOOST_TEST(outputShape[3] == 4);
  }
  
  BOOST_AUTO_TEST_CASE(ExpandDimsInvalidAxisTest)
  {
      armnn::TensorShape inputShape({ 2, 3, 4 });
  
      // Invalid expand dimension 4
      BOOST_CHECK_THROW(ExpandDims(inputShape, 4), armnn::InvalidArgumentException);
  }
  
  BOOST_AUTO_TEST_CASE(ExpandDimsInvalidNegativeAxisTest)
  {
      armnn::TensorShape inputShape({ 2, 3, 4 });
  
      // Invalid expand dimension -5
      BOOST_CHECK_THROW(ExpandDims(inputShape, -5), armnn::InvalidArgumentException);
  }
  
  BOOST_AUTO_TEST_SUITE_END()
  