| 56 | } |
| 57 | |
| 58 | void NEArgMinMaxLayer::configure(ITensor *input, int axis, ITensor *output, const ReductionOperation &op) |
| 59 | { |
| 60 | ARM_COMPUTE_TRACE_EVENT(ARM_COMPUTE_PROF_CAT_CPU, ARM_COMPUTE_PROF_LVL_CPU, "NEArgMinMaxLayer::configure"); |
| 61 | ARM_COMPUTE_LOG_PARAMS(input, axis, output, op); |
| 62 | _impl->reduction_function = std::make_unique<NEReductionOperation>(); |
| 63 | if (output->info() && |
| 64 | (output->info()->data_type() == DataType::S64 || output->info()->data_type() == DataType::U64)) |
| 65 | { |
| 66 | _impl->memory_group = MemoryGroup(std::move(_impl->memory_manager)); |
| 67 | _impl->cast_function = std::make_unique<NECast>(); |
| 68 | _impl->tmp_reduction_result = std::make_unique<Tensor>(); |
| 69 | _impl->reduction_function->configure(input, _impl->tmp_reduction_result.get(), axis, op, false); |
| 70 | _impl->cast_function->configure(_impl->tmp_reduction_result.get(), output, ConvertPolicy::SATURATE); |
| 71 | _impl->memory_group.manage(_impl->tmp_reduction_result.get()); |
| 72 | _impl->tmp_reduction_result->allocator()->allocate(); |
| 73 | } |
| 74 | else |
| 75 | { |
| 76 | _impl->reduction_function->configure(input, output, axis, op, false); |
| 77 | } |
| 78 | } |
| 79 | |
| 80 | Status |
| 81 | NEArgMinMaxLayer::validate(const ITensorInfo *input, int axis, const ITensorInfo *output, const ReductionOperation &op) |