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/add_sub_grad_infer.h"
18 #include "nnacl/arithmetic_parameter.h"
19 #include "nnacl/infer/infer_register.h"
20
AddSubGradInferShape(const TensorC * const * inputs,size_t inputs_size,TensorC ** outputs,size_t outputs_size,OpParameter * parameter)21 int AddSubGradInferShape(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 *dy = inputs[0];
29 const TensorC *x1 = inputs[1];
30 const TensorC *x2 = 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 ArithmeticParameter *param = (ArithmeticParameter *)parameter;
39
40 param->ndim_ = dy->shape_size_;
41 param->in_elements_num0_ = (int)param->ndim_;
42 param->in_elements_num1_ = (int)param->ndim_;
43 param->out_elements_num_ = (int)param->ndim_;
44 size_t fillDimNum0 = dy->shape_size_ - x1->shape_size_;
45 size_t fillDimNum1 = dy->shape_size_ - x2->shape_size_;
46 size_t j0 = 0;
47 size_t j1 = 0;
48 for (size_t i = 0; i < dy->shape_size_; i++) {
49 param->in_shape0_[i] = (i < fillDimNum0) ? 1 : x1->shape_[j0++];
50 param->in_shape1_[i] = (i < fillDimNum1) ? 1 : x2->shape_[j1++];
51 param->out_shape_[i] = dy->shape_[i];
52 }
53
54 SetShapeTensor(dx1, x1);
55 SetShapeTensor(dx2, x2);
56 SetDataTypeFormat(dx1, dy);
57 SetDataTypeFormat(dx2, dy);
58 return NNACL_OK;
59 }
60
61 REG_INFER(AddGrad, PrimType_AddGrad, AddSubGradInferShape)
62 REG_INFER(SubGrad, PrimType_SubGrad, AddSubGradInferShape)
63