diff options
Diffstat (limited to 'src/armnn')
-rw-r--r-- | src/armnn/test/CreateWorkload.hpp | 33 |
1 files changed, 33 insertions, 0 deletions
diff --git a/src/armnn/test/CreateWorkload.hpp b/src/armnn/test/CreateWorkload.hpp index 5308a1c1dc..c0d2ab1c7f 100644 --- a/src/armnn/test/CreateWorkload.hpp +++ b/src/armnn/test/CreateWorkload.hpp @@ -1021,4 +1021,37 @@ std::unique_ptr<ConvertFp32ToFp16Float16Workload> CreateConvertFp32ToFp16Workloa return workload; } +template <typename MeanWorkload, armnn::DataType DataType> +std::unique_ptr<MeanWorkload> CreateMeanWorkloadTest(armnn::IWorkloadFactory& factory, armnn::Graph& graph) +{ + // Reduce along the first and second dimensions, and do not keep the reduced dimensions. + MeanDescriptor descriptor({ 1, 2 }, false); + + // Creates the layer we're testing. + Layer* const layer = graph.AddLayer<MeanLayer>(descriptor, "mean"); + + // Creates extra layers. + Layer* const input = graph.AddLayer<InputLayer>(0, "input"); + Layer* const output = graph.AddLayer<OutputLayer>(0, "output"); + + // Connects up. + armnn::TensorInfo inputTensorInfo({ 1, 3, 7, 4 }, DataType); + armnn::TensorInfo outputTensorInfo({ 1, 4 }, DataType); + Connect(input, layer, inputTensorInfo); + Connect(layer, output, outputTensorInfo); + CreateTensorHandles(graph, factory); + + // Makes the workload and checks it. + auto workload = MakeAndCheckWorkload<MeanWorkload>(*layer, graph, factory); + + MeanQueueDescriptor queueDescriptor = workload->GetData(); + BOOST_TEST(queueDescriptor.m_Parameters.m_Axis == descriptor.m_Axis); + BOOST_TEST(queueDescriptor.m_Parameters.m_KeepDims == descriptor.m_KeepDims); + BOOST_TEST(queueDescriptor.m_Inputs.size() == 1); + BOOST_TEST(queueDescriptor.m_Outputs.size() == 1); + + // Returns so we can do extra, backend-specific tests. + return workload; +} + } |