1 /**
2 * Copyright 2021 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "nnacl/infer/max_min_grad_infer.h"
18 #include "nnacl/arithmetic_parameter.h"
19 #include "nnacl/infer/infer_register.h"
20
MaxMinGradInferShape(const TensorC * const * inputs,size_t inputs_size,TensorC ** outputs,size_t outputs_size,OpParameter * parameter)21 int MaxMinGradInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC **outputs, size_t outputs_size,
22 OpParameter *parameter) {
23 int check_ret = CheckAugmentNullSize(inputs, inputs_size, outputs, outputs_size, parameter, 3, 2);
24 if (check_ret != NNACL_OK) {
25 return check_ret;
26 }
27
28 const TensorC *x1 = inputs[0];
29 const TensorC *x2 = inputs[1];
30 const TensorC *dy = inputs[2];
31 TensorC *dx1 = outputs[0];
32 TensorC *dx2 = outputs[1];
33
34 if (!InferFlag(inputs, inputs_size)) {
35 return NNACL_INFER_INVALID;
36 }
37
38 if (x1->shape_size_ > MAX_SHAPE_SIZE || x2->shape_size_ > MAX_SHAPE_SIZE || dy->shape_size_ > MAX_SHAPE_SIZE) {
39 return NNACL_INPUT_TENSOR_ERROR;
40 }
41 ArithmeticParameter *param = (ArithmeticParameter *)parameter;
42
43 param->ndim_ = dy->shape_size_;
44 param->in_elements_num0_ = (int)(param->ndim_);
45 param->in_elements_num1_ = (int)(param->ndim_);
46 param->out_elements_num_ = (int)(param->ndim_);
47 int fillDimNum0 = (int)(dy->shape_size_ - x1->shape_size_);
48 int fillDimNum1 = (int)(dy->shape_size_ - x2->shape_size_);
49 int j0 = 0;
50 int j1 = 0;
51 for (unsigned int i = 0; i < dy->shape_size_; i++) {
52 param->in_shape0_[i] = ((int)i < fillDimNum0) ? 1 : x1->shape_[j0++];
53 param->in_shape1_[i] = ((int)i < fillDimNum1) ? 1 : x2->shape_[j1++];
54 param->out_shape_[i] = dy->shape_[i];
55 }
56
57 SetShapeTensor(dx1, x1);
58 SetShapeTensor(dx2, x2);
59 SetDataTypeFormat(dx1, dy);
60 SetDataTypeFormat(dx2, dy);
61 return NNACL_OK;
62 }
63
64 REG_INFER(MaximumGrad, PrimType_MaximumGrad, MaxMinGradInferShape)
65 REG_INFER(MinimumGrad, PrimType_MinimumGrad, MaxMinGradInferShape)
66