1 /**
2 * Copyright 2022 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/fused_ada_factor.h"
18
19 #include <memory>
20 #include <vector>
21
22 #include "abstract/abstract_value.h"
23 #include "abstract/dshape.h"
24 #include "abstract/ops/op_infer.h"
25 #include "abstract/ops/primitive_infer_map.h"
26 #include "abstract/utils.h"
27 #include "base/base.h"
28 #include "ir/anf.h"
29 #include "ir/dtype/container.h"
30 #include "ir/primitive.h"
31 #include "mindapi/base/shared_ptr.h"
32 #include "mindapi/ir/value.h"
33 #include "mindapi/src/helper.h"
34 #include "mindspore/core/ops/nn_optimizer_ops.h"
35 #include "ops/primitive_c.h"
36 #include "utils/check_convert_utils.h"
37 #include "utils/convert_utils_base.h"
38 #include "utils/log_adapter.h"
39
40 namespace mindspore {
41 namespace ops {
42 namespace {
43 constexpr size_t kParamIndex = 7;
44 constexpr size_t kFusedAdaFactorInputsNum = 12;
45 auto constexpr kEnableScaleParameter = "enable_scale_parameter";
46 auto constexpr kEnableFirstMoment = "enable_first_moment";
47 auto constexpr kEnableWeightDecay = "enable_weight_decay";
FusedAdaFactorInferShape(const PrimitivePtr & primitive,const std::vector<AbstractBasePtr> & input_args)48 abstract::TupleShapePtr FusedAdaFactorInferShape(const PrimitivePtr &primitive,
49 const std::vector<AbstractBasePtr> &input_args) {
50 MS_EXCEPTION_IF_NULL(primitive);
51 auto param_shape_r = input_args[kParamIndex]->Broaden()->GetShape();
52 auto outputs = std::make_shared<abstract::TupleShape>(std::vector<abstract::BaseShapePtr>({param_shape_r}));
53 return outputs;
54 }
55
FusedAdaFactorInferType(const PrimitivePtr & primitive,const std::vector<AbstractBasePtr> & input_args)56 TypePtr FusedAdaFactorInferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
57 MS_EXCEPTION_IF_NULL(primitive);
58 auto type = input_args[kParamIndex]->GetType();
59 return std::make_shared<Tuple>(std::vector<TypePtr>{type});
60 }
61 } // namespace
62
set_enable_scale_parameter(bool flag)63 void FusedAdaFactor::set_enable_scale_parameter(bool flag) {
64 (void)this->AddAttr(kEnableScaleParameter, api::MakeValue(flag));
65 }
66
get_enable_scale_parameter() const67 bool FusedAdaFactor::get_enable_scale_parameter() const {
68 auto value_ptr = GetAttr(kEnableScaleParameter);
69 return GetValue<bool>(value_ptr);
70 }
71
set_enable_first_moment(bool flag)72 void FusedAdaFactor::set_enable_first_moment(bool flag) {
73 (void)this->AddAttr(kEnableFirstMoment, api::MakeValue(flag));
74 }
75
get_enable_first_moment() const76 bool FusedAdaFactor::get_enable_first_moment() const {
77 auto value_ptr = GetAttr(kEnableFirstMoment);
78 return GetValue<bool>(value_ptr);
79 }
80
set_enable_weight_decay(bool flag)81 void FusedAdaFactor::set_enable_weight_decay(bool flag) {
82 (void)this->AddAttr(kEnableWeightDecay, api::MakeValue(flag));
83 }
84
get_enable_weight_decay() const85 bool FusedAdaFactor::get_enable_weight_decay() const {
86 auto value_ptr = GetAttr(kEnableWeightDecay);
87 return GetValue<bool>(value_ptr);
88 }
89
90 MIND_API_OPERATOR_IMPL(FusedAdaFactor, BaseOperator);
91 MIND_API_OPERATOR_IMPL(FusedAdaFactorWithGlobalNorm, FusedAdaFactor);
FusedAdaFactorInfer(const abstract::AnalysisEnginePtr &,const PrimitivePtr & primitive,const std::vector<AbstractBasePtr> & input_args)92 AbstractBasePtr FusedAdaFactorInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
93 const std::vector<AbstractBasePtr> &input_args) {
94 MS_EXCEPTION_IF_NULL(primitive);
95 for (auto &item : input_args) {
96 MS_EXCEPTION_IF_NULL(item);
97 }
98 auto op_name = primitive->name();
99 (void)CheckAndConvertUtils::CheckInteger("input numbers", SizeToLong(input_args.size()), kGreaterEqual,
100 SizeToLong(kFusedAdaFactorInputsNum), op_name);
101 auto types = FusedAdaFactorInferType(primitive, input_args);
102 auto shapes = FusedAdaFactorInferShape(primitive, input_args);
103 return abstract::MakeAbstract(shapes, types);
104 }
105
106 // AG means auto generated
107 class MIND_API AGFusedAdaFactorInfer : public abstract::OpInferBase {
108 public:
InferShape(const PrimitivePtr & primitive,const std::vector<AbstractBasePtr> & input_args) const109 BaseShapePtr InferShape(const PrimitivePtr &primitive,
110 const std::vector<AbstractBasePtr> &input_args) const override {
111 return input_args[kParamIndex]->GetShape()->Clone();
112 }
113
InferType(const PrimitivePtr & primitive,const std::vector<AbstractBasePtr> & input_args) const114 TypePtr InferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) const override {
115 return FusedAdaFactorInferType(primitive, input_args);
116 }
InferShapeAndType(const abstract::AnalysisEnginePtr & engine,const PrimitivePtr & primitive,const std::vector<AbstractBasePtr> & input_args) const117 AbstractBasePtr InferShapeAndType(const abstract::AnalysisEnginePtr &engine, const PrimitivePtr &primitive,
118 const std::vector<AbstractBasePtr> &input_args) const override {
119 return FusedAdaFactorInfer(engine, primitive, input_args);
120 }
121 };
122
123 REGISTER_PRIMITIVE_OP_INFER_IMPL(FusedAdaFactor, prim::kPrimFusedAdaFactor, AGFusedAdaFactorInfer, false);
124 REGISTER_PRIMITIVE_OP_INFER_IMPL(FusedAdaFactorWithGlobalNorm, prim::kPrimFusedAdaFactorWithGlobalNorm,
125 AGFusedAdaFactorInfer, false);
126 } // namespace ops
127 } // namespace mindspore
128