5 files changed, 70 insertions, 2 deletions
diff --git a/src/armnnSerializer/ArmnnSchema.fbs b/src/armnnSerializer/ArmnnSchema.fbs
index 8b275b6f17..e8d72fc997 100644
--- a/src/armnnSerializer/ArmnnSchema.fbs
+++ b/src/armnnSerializer/ArmnnSchema.fbs
@@ -119,7 +119,8 @@ enum LayerType : uint {
     Lstm = 34,
     Quantize = 35,
     Dequantize = 36,
-    Merge = 37
+    Merge = 37,
+    Switch = 38
 }
 
 // Base layer table to be used as part of other layers
@@ -529,6 +530,10 @@ table MergeLayer {
     base:LayerBase;
 }
 
+table SwitchLayer {
+    base:LayerBase;
+}
+
 union Layer {
     ActivationLayer,
     AdditionLayer,
@@ -567,7 +572,8 @@ union Layer {
     LstmLayer,
     QuantizeLayer,
     DequantizeLayer,
-    MergeLayer
+    MergeLayer,
+    SwitchLayer
 }
 
 table AnyLayer {
diff --git a/src/armnnSerializer/Serializer.cpp b/src/armnnSerializer/Serializer.cpp
index fe30c3eee5..74d0c435c6 100644
--- a/src/armnnSerializer/Serializer.cpp
+++ b/src/armnnSerializer/Serializer.cpp
@@ -893,6 +893,14 @@ void SerializerVisitor::VisitSubtractionLayer(const armnn::IConnectableLayer* la
     CreateAnyLayer(fbSubtractionLayer.o, serializer::Layer::Layer_SubtractionLayer);
 }
 
+void SerializerVisitor::VisitSwitchLayer(const armnn::IConnectableLayer* layer, const char* name)
+{
+    auto fbSwitchBaseLayer = CreateLayerBase(layer, serializer::LayerType::LayerType_Switch);
+    auto fbSwitchLayer = serializer::CreateSwitchLayer(m_flatBufferBuilder, fbSwitchBaseLayer);
+
+    CreateAnyLayer(fbSwitchLayer.o, serializer::Layer::Layer_SwitchLayer);
+}
+
 fb::Offset<serializer::LayerBase> SerializerVisitor::CreateLayerBase(const IConnectableLayer* layer,
                                                                      const serializer::LayerType layerType)
 {
diff --git a/src/armnnSerializer/Serializer.hpp b/src/armnnSerializer/Serializer.hpp
index 775df83966..4a718378b5 100644
--- a/src/armnnSerializer/Serializer.hpp
+++ b/src/armnnSerializer/Serializer.hpp
@@ -191,6 +191,9 @@ public:
 
     void VisitSubtractionLayer(const armnn::IConnectableLayer* layer,
                                const char* name = nullptr) override;
+
+    void VisitSwitchLayer(const armnn::IConnectableLayer* layer,
+                          const char* name = nullptr) override;
 private:
 
     /// Creates the Input Slots and Output Slots and LayerBase for the layer.
diff --git a/src/armnnSerializer/SerializerSupport.md b/src/armnnSerializer/SerializerSupport.md
index a8335e1e68..5b54bfd7be 100644
--- a/src/armnnSerializer/SerializerSupport.md
+++ b/src/armnnSerializer/SerializerSupport.md
@@ -41,5 +41,6 @@ The Arm NN SDK Serializer currently supports the following layers:
 * Splitter
 * StridedSlice
 * Subtraction
+* Switch
 
 More machine learning layers will be supported in future releases.
diff --git a/src/armnnSerializer/test/SerializerTests.cpp b/src/armnnSerializer/test/SerializerTests.cpp
index a1ef9eef59..2724ba4d35 100644
--- a/src/armnnSerializer/test/SerializerTests.cpp
+++ b/src/armnnSerializer/test/SerializerTests.cpp
@@ -2113,6 +2113,56 @@ BOOST_AUTO_TEST_CASE(SerializeSubtraction)
     deserializedNetwork->Accept(verifier);
 }
 
+BOOST_AUTO_TEST_CASE(SerializeSwitch)
+{
+    class SwitchLayerVerifier : public LayerVerifierBase
+    {
+    public:
+        SwitchLayerVerifier(const std::string& layerName,
+                                 const std::vector<armnn::TensorInfo>& inputInfos,
+                                 const std::vector<armnn::TensorInfo>& outputInfos)
+            : LayerVerifierBase(layerName, inputInfos, outputInfos) {}
+
+        void VisitSwitchLayer(const armnn::IConnectableLayer* layer, const char* name) override
+        {
+            VerifyNameAndConnections(layer, name);
+        }
+
+        void VisitConstantLayer(const armnn::IConnectableLayer* layer,
+                                const armnn::ConstTensor& input,
+                                const char *name) override {}
+    };
+
+    const std::string layerName("switch");
+    const armnn::TensorInfo info({ 1, 4 }, armnn::DataType::Float32);
+
+    std::vector<float> constantData = GenerateRandomData<float>(info.GetNumElements());
+    armnn::ConstTensor constTensor(info, constantData);
+
+    armnn::INetworkPtr network = armnn::INetwork::Create();
+    armnn::IConnectableLayer* const inputLayer = network->AddInputLayer(0);
+    armnn::IConnectableLayer* const constantLayer = network->AddConstantLayer(constTensor, "constant");
+    armnn::IConnectableLayer* const switchLayer = network->AddSwitchLayer(layerName.c_str());
+    armnn::IConnectableLayer* const trueOutputLayer = network->AddOutputLayer(0);
+    armnn::IConnectableLayer* const falseOutputLayer = network->AddOutputLayer(1);
+
+    inputLayer->GetOutputSlot(0).Connect(switchLayer->GetInputSlot(0));
+    constantLayer->GetOutputSlot(0).Connect(switchLayer->GetInputSlot(1));
+    switchLayer->GetOutputSlot(0).Connect(trueOutputLayer->GetInputSlot(0));
+    switchLayer->GetOutputSlot(1).Connect(falseOutputLayer->GetInputSlot(0));
+
+    inputLayer->GetOutputSlot(0).SetTensorInfo(info);
+    constantLayer->GetOutputSlot(0).SetTensorInfo(info);
+    switchLayer->GetOutputSlot(0).SetTensorInfo(info);
+    switchLayer->GetOutputSlot(1).SetTensorInfo(info);
+
+    armnn::INetworkPtr deserializedNetwork = DeserializeNetwork(SerializeNetwork(*network));
+    BOOST_CHECK(deserializedNetwork);
+
+    SwitchLayerVerifier verifier(layerName, {info, info}, {info, info});
+    deserializedNetwork->Accept(verifier);
+}
+
 BOOST_AUTO_TEST_CASE(SerializeDeserializeNonLinearNetwork)
 {
     class ConstantLayerVerifier : public LayerVerifierBase