diff options
Diffstat (limited to 'src/runtime/NEON/functions/NEReduceMean.cpp')
| -rw-r--r-- | src/runtime/NEON/functions/NEReduceMean.cpp | 12 |
1 files changed, 10 insertions, 2 deletions
diff --git a/src/runtime/NEON/functions/NEReduceMean.cpp b/src/runtime/NEON/functions/NEReduceMean.cpp index 4547a1f9b0..96ec8b8587 100644 --- a/src/runtime/NEON/functions/NEReduceMean.cpp +++ b/src/runtime/NEON/functions/NEReduceMean.cpp @@ -26,9 +26,11 @@ #include "arm_compute/core/CPP/Validate.h" #include "arm_compute/core/Error.h" #include "arm_compute/core/Helpers.h" +#include "arm_compute/core/utils/misc/ShapeCalculator.h" #include "arm_compute/runtime/NEON/NEScheduler.h" -using namespace arm_compute; +namespace arm_compute +{ NEReduceMean::NEReduceMean(std::shared_ptr<IMemoryManager> memory_manager) : _memory_group(std::move(memory_manager)), _reduction_kernels(), _reduced_outs(), _reshape(), _reduction_ops(), _keep_dims() @@ -95,7 +97,11 @@ Status NEReduceMean::validate(const ITensorInfo *input, const Coordinates &reduc void NEReduceMean::configure(ITensor *input, const Coordinates &reduction_axis, bool keep_dims, ITensor *output) { - ARM_COMPUTE_ERROR_ON_NULLPTR(input); + // Perform validate step + ARM_COMPUTE_ERROR_THROW_ON(NEReduceMean::validate(input->info(), reduction_axis, keep_dims, output->info())); + // Output auto inizialitation if not yet initialized + const TensorShape output_shape = arm_compute::misc::shape_calculator::calculate_reduce_mean_shape(input, reduction_axis, keep_dims); + auto_init_if_empty(*output->info(), input->info()->clone()->set_tensor_shape(output_shape)); _reduction_ops = reduction_axis.num_dimensions(); _reduction_kernels.resize(_reduction_ops); @@ -161,3 +167,5 @@ void NEReduceMean::run() _reshape.run(); } } + +} // namespace arm_compute |