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/constant_of_shape_infer.h"
18 #include "nnacl/infer/infer_register.h"
19 #include "nnacl/tensor_c_utils.h"
20
ConstantOfShapeInferShape(const TensorC * const * inputs,size_t inputs_size,TensorC ** outputs,size_t outputs_size,OpParameter * parameter)21 int ConstantOfShapeInferShape(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, 1, 1);
24 if (check_ret != NNACL_OK) {
25 return check_ret;
26 }
27
28 const TensorC *in_tensor = inputs[0];
29 TensorC *out_tensor = outputs[0];
30 ConstantOfShapeParameter *param = (ConstantOfShapeParameter *)parameter;
31 out_tensor->data_type_ = (TypeIdC)(param->data_type_);
32 out_tensor->format_ = in_tensor->format_;
33 if (!InferFlag(inputs, inputs_size) || in_tensor->data_ == NULL) {
34 return NNACL_INFER_INVALID;
35 }
36 int size = GetElementNum(in_tensor);
37 if (size < 0 || size > MAX_SHAPE_SIZE) {
38 return NNACL_ERR;
39 }
40 int out_shape[MAX_SHAPE_SIZE];
41 int out_shape_size = size;
42 switch (in_tensor->data_type_) {
43 case kNumberTypeInt32: {
44 int32_t *in_data = (int32_t *)(in_tensor->data_);
45 for (int i = 0; i < size; ++i) {
46 out_shape[i] = in_data[i];
47 if (out_shape[i] < 0) {
48 return NNACL_ERR;
49 }
50 }
51 break;
52 }
53 case kNumberTypeInt64: {
54 int64_t *in_data = (int64_t *)(in_tensor->data_);
55 for (int i = 0; i < size; ++i) {
56 out_shape[i] = in_data[i];
57 if (out_shape[i] < 0) {
58 return NNACL_ERR;
59 }
60 }
61 break;
62 }
63 default:
64 return NNACL_INFER_INVALID;
65 }
66
67 SetShapeArray(out_tensor, out_shape, (size_t)out_shape_size);
68 return NNACL_OK;
69 }
70
71 REG_INFER(ConstantOfShape, PrimType_ConstantOfShape, ConstantOfShapeInferShape)
72