1 files changed, 29 insertions, 0 deletions
diff --git a/src/armnn/test/QuantizerTest.cpp b/src/armnn/test/QuantizerTest.cpp
index 548203a6a9..fcce208c59 100644
--- a/src/armnn/test/QuantizerTest.cpp
+++ b/src/armnn/test/QuantizerTest.cpp
@@ -14,6 +14,7 @@
 #include "../NetworkQuantizerUtils.hpp"
 #include "../OverrideInputRangeVisitor.hpp"
 #include "../RangeTracker.hpp"
+#include "../backends/backendsCommon/test/QuantizeHelper.hpp"
 
 #include <boost/test/unit_test.hpp>
 
@@ -1215,5 +1216,33 @@ BOOST_AUTO_TEST_CASE(QuantizeBatchToSpace)
     VisitLayersTopologically(quantizedNetwork.get(), validator);
 }
 
+std::vector<uint8_t> SetupQuantize(float value)
+{
+    armnn::TensorInfo inputInfo({ 1, 2, 2 }, armnn::DataType::Float32);
+    inputInfo.SetQuantizationScale(1.0f);
+    inputInfo.SetQuantizationOffset(1);
+    std::vector<float> input({
+                                     value, 0.0f,
+                                     0.0f, 1.0f
+                             });
+    const std::vector<float> &inputRef = input;
+
+    auto output = QuantizedVector<uint8_t>(inputInfo.GetQuantizationScale(),
+                                           inputInfo.GetQuantizationOffset(),
+                                           inputRef);
+
+    return output;
+}
+
+BOOST_AUTO_TEST_CASE(QuantizeInf)
+{
+    BOOST_CHECK_EQUAL(SetupQuantize(std::numeric_limits<float>::infinity())[0], 255);
+}
+
+BOOST_AUTO_TEST_CASE(QuantizeNegativeInf)
+{
+    BOOST_CHECK_EQUAL(SetupQuantize(-1 * std::numeric_limits<float>::infinity())[0], 0);
+}
+
 BOOST_AUTO_TEST_SUITE_END()
 } // namespace armnn