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 <map>
18 #include <string>
19 #include "ops/concat.h"
20 #include "ops/op_utils.h"
21 #include "utils/check_convert_utils.h"
22 namespace mindspore {
23 namespace ops {
Init(const int64_t axis)24 void Concat::Init(const int64_t axis) { this->set_axis(axis); }
get_axis() const25 int64_t Concat::get_axis() const {
26 auto value_ptr = this->GetAttr(kAxis);
27 return GetValue<int64_t>(value_ptr);
28 }
29
set_axis(const int64_t axis)30 void Concat::set_axis(const int64_t axis) { (void)this->AddAttr(kAxis, MakeValue(axis)); }
31
ConcatInfer(const abstract::AnalysisEnginePtr &,const PrimitivePtr & primitive,const std::vector<AbstractBasePtr> & input_args)32 AbstractBasePtr ConcatInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
33 const std::vector<AbstractBasePtr> &input_args) {
34 MS_EXCEPTION_IF_NULL(primitive);
35 auto prim_name = primitive->name();
36 (void)CheckAndConvertUtils::CheckInteger("input number", SizeToLong(input_args.size()), kEqual, 1, prim_name);
37 for (const auto &item : input_args) {
38 MS_EXCEPTION_IF_NULL(item);
39 }
40 auto input_tuple = input_args[0]->cast<abstract::AbstractTuplePtr>();
41 MS_EXCEPTION_IF_NULL(input_tuple);
42 auto elements = input_tuple->elements();
43 const int64_t kOneNum = 1;
44 (void)CheckAndConvertUtils::CheckInteger("concat element num", SizeToLong(elements.size()), kGreaterEqual, kOneNum,
45 prim_name);
46 auto element0 = elements[0]->cast<abstract::AbstractTensorPtr>();
47 MS_EXCEPTION_IF_NULL(element0);
48 auto element0_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(element0->BuildShape())[kShape];
49 auto element0_rank = element0_shape.size();
50 auto axis_temp = GetValue<int64_t>(primitive->GetAttr(kAxis));
51 CheckAndConvertUtils::CheckInRange<int64_t>("Concat axis", axis_temp, kIncludeBoth,
52 {-SizeToLong(element0_rank) - kOneNum, SizeToLong(element0_rank)},
53 prim_name);
54 auto axis = axis_temp < 0 ? LongToSize(axis_temp) + element0_rank : LongToSize(axis_temp);
55
56 std::map<std::string, TypePtr> types;
57 (void)types.emplace("element0", element0->BuildType());
58 int64_t all_shp = element0_shape[axis];
59 for (size_t i = 1; i < elements.size(); ++i) {
60 std::string elementi = "element" + std::to_string(i);
61 auto elementi_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(elements[i]->BuildShape())[kShape];
62 (void)CheckAndConvertUtils::CheckInteger(elementi + " shape rank", SizeToLong(elementi_shape.size()), kEqual,
63 SizeToLong(element0_shape.size()), prim_name);
64 for (size_t j = 0; j < element0_rank; ++j) {
65 if (j != axis && elementi_shape[j] != element0_shape[j]) {
66 MS_LOG(EXCEPTION) << "element " << i << " shape in input can not concat with first element.";
67 }
68 }
69 all_shp = all_shp == -1 || elementi_shape[axis] == -1 ? -1 : all_shp + elementi_shape[axis];
70 (void)types.emplace(elementi, elements[i]->BuildType());
71 }
72 auto infer_type = CheckAndConvertUtils::CheckTensorTypeSame(types, all_types, prim_name);
73 auto ret_shape = element0_shape;
74 ret_shape[axis] = all_shp;
75 return std::make_shared<abstract::AbstractTensor>(infer_type, std::make_shared<abstract::Shape>(ret_shape));
76 }
77 REGISTER_PRIMITIVE_C(kNameConcat, Concat);
78 } // namespace ops
79 } // namespace mindspore
80