1 /**
2 * Copyright 2020 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 <vector>
18 #include <functional>
19
20 #include "ops/tensor_list_stack.h"
21 #include "ops/op_utils.h"
22 #include "utils/check_convert_utils.h"
23
24 namespace mindspore {
25 namespace ops {
Init(const int64_t num_elements,const int64_t element_dtype)26 void TensorListStack::Init(const int64_t num_elements, const int64_t element_dtype) {
27 this->set_num_elements(num_elements);
28 this->set_element_dtype(element_dtype);
29 }
30
set_num_elements(const int64_t num_elements)31 void TensorListStack::set_num_elements(const int64_t num_elements) {
32 (void)this->AddAttr(kNumElements, MakeValue(num_elements));
33 }
34
set_element_dtype(const int64_t element_dtype)35 void TensorListStack::set_element_dtype(const int64_t element_dtype) {
36 (void)this->AddAttr(kElement_dtype, MakeValue(element_dtype));
37 }
38
get_num_elements() const39 int64_t TensorListStack::get_num_elements() const {
40 auto value_ptr = GetAttr(kNumElements);
41 return GetValue<int64_t>(value_ptr);
42 }
43
get_element_dtype() const44 int64_t TensorListStack::get_element_dtype() const {
45 auto value_ptr = GetAttr(kElement_dtype);
46 return GetValue<int64_t>(value_ptr);
47 }
48
TensorListStackInfer(const abstract::AnalysisEnginePtr &,const PrimitivePtr & primitive,const std::vector<AbstractBasePtr> & input_args)49 AbstractBasePtr TensorListStackInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
50 const std::vector<AbstractBasePtr> &input_args) {
51 MS_EXCEPTION_IF_NULL(primitive);
52 for (const auto &input : input_args) {
53 MS_EXCEPTION_IF_NULL(input);
54 }
55 auto input0_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->BuildShape())[kShape];
56 int64_t num = std::accumulate(input0_shape.begin(), input0_shape.end(), 1LL, std::multiplies<int64_t>());
57 if (num == 0) {
58 MS_LOG(ERROR) << "Try to stack a empty tensorlist!";
59 }
60 if (input_args[1]->BuildShape() == nullptr) {
61 MS_LOG(ERROR) << "ele_shape->data_c() is nullptr";
62 }
63 auto input1_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[1]->BuildShape())[kShape];
64 (void)input1_shape.insert(input1_shape.begin(), 1);
65 return std::make_shared<abstract::AbstractTensor>(input_args[0]->BuildType(), input1_shape);
66 }
67 REGISTER_PRIMITIVE_C(kNameTensorListStack, TensorListStack);
68 } // namespace ops
69 } // namespace mindspore
70