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/pack.h"
18
19 namespace mindspore {
20 namespace ops {
21 namespace {
_get_pack_shape(std::vector<BaseShapePtr> x_shapes,std::vector<TypePtr> x_types,int64_t axis,const std::string & name)22 std::vector<int64_t> _get_pack_shape(std::vector<BaseShapePtr> x_shapes, std::vector<TypePtr> x_types, int64_t axis,
23 const std::string &name) {
24 (void)CheckAndConvertUtils::CheckInteger("len of input_x", (int64_t)x_shapes.size(), kGreaterEqual, 1, name);
25 (void)CheckAndConvertUtils::CheckSubClass("input_x[0]", x_types[0], {TypeIdToType(kObjectTypeTensorType)}, name);
26 auto output_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(x_shapes[0])[kShape];
27 int64_t rank_base = SizeToLong(output_shape.size());
28 int64_t N = SizeToLong(x_shapes.size());
29 if (axis < 0) {
30 axis = axis + rank_base + 1;
31 }
32 for (int64_t i = 1; i < N; i++) {
33 auto type = x_types[LongToSize(i)]->cast<TensorTypePtr>()->element();
34 MS_EXCEPTION_IF_NULL(type);
35 auto type0 = x_types[0]->cast<TensorTypePtr>()->element();
36 MS_EXCEPTION_IF_NULL(type0);
37 CheckAndConvertUtils::Check("x_type[" + std::to_string(i) + "]", type->type_id(), kEqual, "base", type0->type_id(),
38 name);
39 auto shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(x_shapes[LongToSize(i)])[kShape];
40 if (shape != output_shape) {
41 MS_EXCEPTION(ValueError) << "For '" + name + "' element " + std::to_string(i) +
42 "shape in input can't pack with first element.";
43 }
44 }
45 (void)output_shape.insert(output_shape.begin() + axis, N);
46 return output_shape;
47 }
48 } // namespace
49
set_axis(const int64_t & axis)50 void Pack::set_axis(const int64_t &axis) { (void)AddAttr(kAxis, MakeValue(axis)); }
51
get_axis() const52 int64_t Pack::get_axis() const { return GetValue<int64_t>(GetAttr(kAxis)); }
53
Init(const int64_t & axis)54 void Pack::Init(const int64_t &axis) { this->set_axis(axis); }
55
PackInfer(const abstract::AnalysisEnginePtr &,const PrimitivePtr & primitive,const std::vector<AbstractBasePtr> & input_args)56 AbstractBasePtr PackInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
57 const std::vector<AbstractBasePtr> &input_args) {
58 MS_EXCEPTION_IF_NULL(primitive);
59 auto prim_name = primitive->name();
60
61 MS_EXCEPTION_IF_NULL(input_args[0]);
62 auto x_shapes = input_args[0]->BuildShape()->cast<abstract::TupleShapePtr>()->shape();
63 auto x_types = input_args[0]->BuildType()->cast<TuplePtr>()->elements();
64 auto all_shape = _get_pack_shape(x_shapes, x_types, GetValue<int64_t>(primitive->GetAttr(kAxis)), prim_name);
65 auto tensor_type = x_types[0]->cast<TensorTypePtr>();
66 MS_EXCEPTION_IF_NULL(tensor_type);
67 auto data_type = tensor_type->element();
68 MS_EXCEPTION_IF_NULL(data_type);
69 return std::make_shared<abstract::AbstractTensor>(data_type, all_shape);
70 }
71 REGISTER_PRIMITIVE_C(kNamePack, Pack);
72 } // namespace ops
73 } // namespace mindspore
74