9 #include <boost/numeric/conversion/cast.hpp> 20 unsigned int carry = 1;
22 for (
unsigned int idx = numDims; idx-- > 0; )
24 unsigned int current_val = current[idx] + carry;
25 if (dims[idx] == current_val)
31 current[idx] = current_val;
41 std::vector<unsigned int>& index,
42 const unsigned int numAxis,
43 const std::vector<unsigned int>& axis)
45 unsigned int offset = 0;
46 for (
unsigned int idx = 0; idx < numDims; ++idx)
51 for (
unsigned int axisIdx = 0; axisIdx < numAxis; ++axisIdx)
53 if (idx == axis[axisIdx])
62 offset = offset * dims[idx] + index[idx];
73 const std::vector<unsigned int>& axis,
85 unsigned int numOutputs = 1;
86 for (
unsigned int idx = 0; idx < outputNumDims; ++idx)
88 numOutputs *= outputDims[idx];
91 std::vector<float> tempSum(numOutputs);
92 for (
unsigned int idx = 0; idx < numOutputs; ++idx)
100 std::vector<unsigned int> tempIndex(inputNumDims);
101 for (
unsigned int idx = 0; idx < inputNumDims; ++idx)
106 std::vector<unsigned int> resolvedAxis = axis;
107 if (resolvedAxis.empty())
109 for (
unsigned int idx = 0; idx < inputNumDims; ++idx)
111 resolvedAxis.push_back(idx);
114 auto numResolvedAxis = boost::numeric_cast<
unsigned int>(resolvedAxis.size());
117 for (
bool hasNext =
true; hasNext; hasNext =
NextIndex(inputNumDims, inputDims, tempIndex))
121 numResolvedAxis, resolvedAxis);
123 tempSum[outputOffset] += input.
Get();
127 size_t numElementsInAxis = 1;
128 for (
unsigned int idx = 0; idx < numResolvedAxis; ++idx)
130 unsigned int current = inputDims[resolvedAxis[idx]];
131 BOOST_ASSERT(boost::numeric_cast<float>(current) <
132 (std::numeric_limits<float>::max() / boost::numeric_cast<float>(numElementsInAxis)));
133 numElementsInAxis *= current;
135 if (numElementsInAxis > 0) {
136 for (
unsigned int idx = 0; idx < numOutputs; ++idx)
139 output.
Set(tempSum[idx] / boost::numeric_cast<float>(numElementsInAxis));
unsigned int GetNumDimensions() const
bool NextIndex(const unsigned int numDims, const armnn::TensorShape &dims, std::vector< unsigned int > ¤t)
void Mean(const armnn::TensorInfo &inputInfo, const armnn::TensorInfo &outputInfo, const std::vector< unsigned int > &axis, Decoder< float > &input, Encoder< float > &output)
unsigned int ReducedOutputOffset(const unsigned int numDims, const armnn::TensorShape &dims, std::vector< unsigned int > &index, const unsigned int numAxis, const std::vector< unsigned int > &axis)
virtual IType Get() const =0
const TensorShape & GetShape() const
virtual void Set(IType right)=0