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 "ops/tensor_list_from_tensor.h"
18 #include "ops/op_utils.h"
19 #include "utils/check_convert_utils.h"
20
21 namespace mindspore {
22 namespace ops {
23 namespace {
TensorListFromTensorInferShape(const PrimitivePtr & primitive,const std::vector<AbstractBasePtr> & input_args)24 abstract::ShapePtr TensorListFromTensorInferShape(const PrimitivePtr &primitive,
25 const std::vector<AbstractBasePtr> &input_args) {
26 MS_EXCEPTION_IF_NULL(primitive);
27 auto op_name = primitive->name();
28 const int64_t input_num = 2;
29 CheckAndConvertUtils::CheckInputArgs(input_args, kGreaterEqual, input_num, op_name);
30
31 auto input0_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->BuildShape())[kShape];
32 if (input0_shape.size() < 1) {
33 MS_LOG(ERROR) << "input0_shape.size():" << input0_shape.size() << " must be greater than 0!";
34 }
35 int64_t dim0 = input0_shape[0];
36 if (dim0 < 0) {
37 MS_LOG(ERROR) << "input[0] dim0:" << dim0 << " must be greater than or equal to 0!";
38 }
39 std::vector<int64_t> infer_shape = {1, dim0};
40 return std::make_shared<abstract::Shape>(infer_shape);
41 }
42
TensorListFromTensorInferType()43 TypePtr TensorListFromTensorInferType() { return kTensorType; }
44 } // namespace
45
Init(const int64_t element_dtype,const int64_t shape_type)46 void TensorListFromTensor::Init(const int64_t element_dtype, const int64_t shape_type) {
47 this->set_element_dtype(element_dtype);
48 this->set_shape_type(shape_type);
49 }
50
get_element_dtype() const51 int64_t TensorListFromTensor::get_element_dtype() const {
52 auto value_ptr = GetAttr(kElement_dtype);
53 return GetValue<int64_t>(value_ptr);
54 }
55
get_shape_type() const56 int64_t TensorListFromTensor::get_shape_type() const {
57 auto value_ptr = GetAttr(kShapeType);
58 return GetValue<int64_t>(value_ptr);
59 }
60
set_element_dtype(const int64_t element_dtype)61 void TensorListFromTensor::set_element_dtype(const int64_t element_dtype) {
62 (void)this->AddAttr(kElement_dtype, MakeValue(element_dtype));
63 }
64
set_shape_type(const int64_t shape_type)65 void TensorListFromTensor::set_shape_type(const int64_t shape_type) {
66 (void)this->AddAttr(kShapeType, MakeValue(shape_type));
67 }
68
TensorListFromTensorInfer(const abstract::AnalysisEnginePtr &,const PrimitivePtr & primitive,const std::vector<AbstractBasePtr> & input_args)69 AbstractBasePtr TensorListFromTensorInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
70 const std::vector<AbstractBasePtr> &input_args) {
71 return std::make_shared<abstract::AbstractTensor>(TensorListFromTensorInferType(),
72 TensorListFromTensorInferShape(primitive, input_args)->shape());
73 }
74 REGISTER_PRIMITIVE_C(kNameTensorListFromTensor, TensorListFromTensor);
75 } // namespace ops
76 } // namespace mindspore
77