1 files changed, 47 insertions, 15 deletions

diff --git a/src/armnnTfLiteParser/test/Split.cpp b/src/armnnTfLiteParser/test/Split.cpp
index 774a416750..a6875143fa 100644
--- a/src/armnnTfLiteParser/test/Split.cpp
+++ b/src/armnnTfLiteParser/test/Split.cpp
@@ -14,11 +14,12 @@ BOOST_AUTO_TEST_SUITE(TensorflowLiteParser)
 
 struct SplitFixture : public ParserFlatbuffersFixture
 {
-    explicit SplitFixture(const std::string & inputShape,
-                          const std::string & axisShape,
-                          const std::string & numSplits,
-                          const std::string & outputShape1,
-                          const std::string & outputShape2)
+    explicit SplitFixture(const std::string& inputShape,
+                          const std::string& axisShape,
+                          const std::string& numSplits,
+                          const std::string& outputShape1,
+                          const std::string& outputShape2,
+                          const std::string& axisData)
     {
         m_JsonString = R"(
             {
@@ -75,12 +76,12 @@ struct SplitFixture : public ParserFlatbuffersFixture
                             }
                         }
                     ],
-                    "inputs": [ 0, 1 ],
+                    "inputs": [ 0 ],
                     "outputs": [ 2, 3 ],
                     "operators": [
                         {
                             "opcode_index": 0,
-                            "inputs": [ 0, 1 ],
+                            "inputs": [ 1, 0 ],
                             "outputs": [ 2, 3 ],
                             "builtin_options_type": "SplitOptions",
                             "builtin_options": {
@@ -90,7 +91,7 @@ struct SplitFixture : public ParserFlatbuffersFixture
                         }
                     ],
                 } ],
-                "buffers" : [ {}, {} ]
+                "buffers" : [ {}, {"data": )" + axisData + R"( }, {}, {} ]
             }
         )";
 
@@ -101,8 +102,8 @@ struct SplitFixture : public ParserFlatbuffersFixture
 
 struct SimpleSplitFixture : SplitFixture
 {
-    SimpleSplitFixture() : SplitFixture( "[ 2, 2, 2, 2 ]", "[ 1 ]", "2",
-        "[ 2, 1, 2, 2 ]", "[ 2, 1, 2, 2 ]")
+    SimpleSplitFixture() : SplitFixture( "[ 2, 2, 2, 2 ]", "[ ]", "2",
+        "[ 2, 1, 2, 2 ]", "[ 2, 1, 2, 2 ]", "[ 1, 0, 0, 0 ]")
          {}
 };
 
@@ -113,14 +114,14 @@ BOOST_FIXTURE_TEST_CASE(ParseAxisOneSplitTwo, SimpleSplitFixture)
         0,
         { {"inputTensor", { 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f,
                             11.0f, 12.0f, 13.0f, 14.0f, 15.0f, 16.0f } } },
-        { {"outputTensor1", { 1.0f, 2.0f, 3.0f, 4.0f, 9.0f, 10.0f, 11.0f, 12.0f }},
-          {"outputTensor2", { 5.0f, 6.0f, 7.0f, 8.0f, 13.0f, 14.0f, 15.0f, 16.0f }}});
+        { {"outputTensor1", { 1.0f, 2.0f, 3.0f, 4.0f, 9.0f, 10.0f, 11.0f, 12.0f } },
+          {"outputTensor2", { 5.0f, 6.0f, 7.0f, 8.0f, 13.0f, 14.0f, 15.0f, 16.0f } } });
 }
 
 struct SimpleSplitAxisThreeFixture : SplitFixture
 {
-    SimpleSplitAxisThreeFixture() : SplitFixture( "[ 2, 2, 2, 2 ]", "[ 3 ]", "2",
-        "[ 2, 2, 2, 1 ]", "[ 2, 2, 2, 1 ]")
+    SimpleSplitAxisThreeFixture() : SplitFixture( "[ 2, 2, 2, 2 ]", "[ ]", "2",
+        "[ 2, 2, 2, 1 ]", "[ 2, 2, 2, 1 ]", "[ 3, 0, 0, 0 ]")
     {}
 };
 
@@ -130,8 +131,39 @@ BOOST_FIXTURE_TEST_CASE(ParseAxisThreeSplitTwo, SimpleSplitAxisThreeFixture)
         0,
         { {"inputTensor", { 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f,
                             11.0f, 12.0f, 13.0f, 14.0f, 15.0f, 16.0f } } },
-        { {"outputTensor1", { 1.0f, 3.0f, 5.0f, 7.0f, 9.0f, 11.0f, 13.0f, 15.0f }},
+        { {"outputTensor1", { 1.0f, 3.0f, 5.0f, 7.0f, 9.0f, 11.0f, 13.0f, 15.0f } },
           {"outputTensor2", { 2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f, 14.0f, 16.0f } } } );
 }
 
+struct SimpleSplit2DFixture : SplitFixture
+{
+    SimpleSplit2DFixture() : SplitFixture( "[ 1, 8 ]", "[ ]", "2", "[ 1, 4 ]", "[ 1, 4 ]", "[ 1, 0, 0, 0 ]")
+    {}
+};
+
+BOOST_FIXTURE_TEST_CASE(SimpleSplit2D, SimpleSplit2DFixture)
+{
+    RunTest<2, armnn::DataType::Float32>(
+            0,
+            { {"inputTensor", { 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f } } },
+            { {"outputTensor1", { 1.0f, 2.0f, 3.0f, 4.0f } },
+              {"outputTensor2", { 5.0f, 6.0f, 7.0f, 8.0f } } } );
+}
+
+struct SimpleSplit3DFixture : SplitFixture
+{
+    SimpleSplit3DFixture() : SplitFixture( "[ 1, 8, 2 ]", "[ ]", "2", "[ 1, 4, 2 ]", "[ 1, 4, 2 ]", "[ 1, 0, 0, 0 ]")
+    {}
+};
+
+BOOST_FIXTURE_TEST_CASE(SimpleSplit3D, SimpleSplit3DFixture)
+{
+    RunTest<3, armnn::DataType::Float32>(
+            0,
+            { {"inputTensor", { 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f,
+                                10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f, 16.0f } } },
+            { {"outputTensor1", { 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f } },
+              {"outputTensor2", { 9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f, 16.0f } } } );
+}
+
 BOOST_AUTO_TEST_SUITE_END()
\ No newline at end of file