1 /*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20 /*!
21 * \file np_init_op.cc
22 * \brief CPU Implementation of numpy init op
23 */
24
25 #include "../tensor/init_op.h"
26 #include "../tensor/elemwise_unary_op.h"
27 #include "./np_init_op.h"
28
29 namespace mxnet {
30 namespace op {
31
32
33 DMLC_REGISTER_PARAMETER(NumpyEyeParam);
34 DMLC_REGISTER_PARAMETER(IndicesOpParam);
35 DMLC_REGISTER_PARAMETER(LogspaceParam);
36 DMLC_REGISTER_PARAMETER(FullLikeOpParam);
37 DMLC_REGISTER_PARAMETER(AtleastNDParam);
38
NumpyIndicesShape(const nnvm::NodeAttrs & attrs,mxnet::ShapeVector * in_shapes,mxnet::ShapeVector * out_shapes)39 inline bool NumpyIndicesShape(const nnvm::NodeAttrs& attrs,
40 mxnet::ShapeVector* in_shapes,
41 mxnet::ShapeVector* out_shapes) {
42 const IndicesOpParam& param = nnvm::get<IndicesOpParam>(attrs.parsed);
43 CHECK_EQ(in_shapes->size(), 0U);
44 CHECK_EQ(out_shapes->size(), 1U);
45 CHECK_GE(param.dimensions.ndim(), 0)
46 << "_npi_indices dimensions the number of dim must not be less than 0";
47 mxnet::TShape param_dim = param.dimensions;
48 if (!shape_is_known(param_dim)) return false;
49 const int indim = param.dimensions.ndim();
50 mxnet::TShape ret(indim + 1, -1);
51 ret[0] = indim;
52 for (int i = 1; i < indim + 1; ++i) {
53 ret[i] = param.dimensions[i-1];
54 }
55 SHAPE_ASSIGN_CHECK(*out_shapes, 0, ret);
56 return shape_is_known(out_shapes->at(0));
57 }
58
LogspaceShape(const nnvm::NodeAttrs & attrs,mxnet::ShapeVector * in_attrs,mxnet::ShapeVector * out_attrs)59 inline bool LogspaceShape(const nnvm::NodeAttrs& attrs,
60 mxnet::ShapeVector *in_attrs,
61 mxnet::ShapeVector *out_attrs) {
62 const LogspaceParam& param = nnvm::get<LogspaceParam>(attrs.parsed);
63 CHECK_EQ(in_attrs->size(), 0U);
64 CHECK_EQ(out_attrs->size(), 1U);
65 CHECK_GE(param.num, 0)
66 << "Number of sequence should be non-negative, received " << param.num;
67 SHAPE_ASSIGN_CHECK(*out_attrs, 0, mxnet::TShape({static_cast<nnvm::dim_t>(param.num)}));
68 return true;
69 }
70
71 NNVM_REGISTER_OP(_npi_zeros)
72 .set_num_inputs(0)
73 .set_num_outputs(1)
74 .set_attr_parser(ParamParser<InitOpParam>)
75 .set_attr<mxnet::FInferShape>("FInferShape", InitShape<InitOpParam>)
76 .set_attr<nnvm::FInferType>("FInferType", InitType<InitOpParam>)
77 .set_attr<FInferStorageType>("FInferStorageType", InitStorageType<InitOpParam, true, true>)
78 .set_attr<FCompute>("FCompute<cpu>", FillCompute<cpu, 0>)
79 .add_arguments(InitOpParam::__FIELDS__());
80
81 NNVM_REGISTER_OP(_npi_ones)
82 .describe("Return a new array of given shape, type, and context, filled with ones.")
83 .set_num_inputs(0)
84 .set_num_outputs(1)
85 .set_attr_parser(ParamParser<InitOpParam>)
86 .set_attr<mxnet::FInferShape>("FInferShape", InitShape<InitOpParam>)
87 .set_attr<nnvm::FInferType>("FInferType", InitType<InitOpParam>)
88 .set_attr<FCompute>("FCompute<cpu>", FillCompute<cpu, 1>)
89 .add_arguments(InitOpParam::__FIELDS__());
90
91 NNVM_REGISTER_OP(_npi_identity)
92 .describe("Return a new identity array of given shape, type, and context.")
93 .set_num_inputs(0)
94 .set_num_outputs(1)
95 .set_attr_parser(ParamParser<InitOpParam>)
96 .set_attr<mxnet::FInferShape>("FInferShape", InitShape<InitOpParam>)
97 .set_attr<nnvm::FInferType>("FInferType", InitType<InitOpParam>)
98 .set_attr<FCompute>("FCompute<cpu>", IdentityCompute<cpu>)
99 .add_arguments(InitOpParam::__FIELDS__());
100
101 template<int NDim>
AtleastNDShape(const nnvm::NodeAttrs & attrs,std::vector<mxnet::TShape> * in_attrs,std::vector<mxnet::TShape> * out_attrs)102 inline bool AtleastNDShape(const nnvm::NodeAttrs& attrs,
103 std::vector<mxnet::TShape> *in_attrs,
104 std::vector<mxnet::TShape> *out_attrs) {
105 auto ¶m = nnvm::get<AtleastNDParam>(attrs.parsed);
106
107 CHECK_EQ(in_attrs->size(), param.num_args);
108 CHECK_EQ(out_attrs->size(), param.num_args);
109
110 for (int i = 0; i < param.num_args; ++i) {
111 auto &shape = in_attrs->at(i);
112 if (shape.ndim() < NDim) {
113 mxnet::TShape new_shape(NDim, 1);
114 if (NDim == 2) {
115 if (shape.ndim() == 1) {
116 new_shape[1] = shape[0];
117 }
118 } else if (NDim == 3) {
119 if (shape.ndim() == 1) {
120 new_shape[1] = shape[0];
121 } else if (shape.ndim() == 2) {
122 new_shape[0] = shape[0];
123 new_shape[1] = shape[1];
124 }
125 }
126 SHAPE_ASSIGN_CHECK(*out_attrs, i, new_shape);
127 } else {
128 SHAPE_ASSIGN_CHECK(*out_attrs, i, shape);
129 }
130 }
131
132 return shape_is_known(*in_attrs) && shape_is_known(*out_attrs);
133 }
134
135 #define NNVM_REGISTER_ATLEAST_ND(N) \
136 NNVM_REGISTER_OP(_np_atleast_##N##d) \
137 .set_attr_parser(ParamParser<AtleastNDParam>) \
138 .set_num_inputs( \
139 [](const NodeAttrs& attrs) { \
140 auto ¶m = nnvm::get<AtleastNDParam>(attrs.parsed); \
141 return param.num_args; \
142 }) \
143 .set_num_outputs( \
144 [](const NodeAttrs& attrs) { \
145 auto ¶m = nnvm::get<AtleastNDParam>(attrs.parsed); \
146 return param.num_args; \
147 }) \
148 .set_attr<std::string>("key_var_num_args", "num_args") \
149 .set_attr<nnvm::FListInputNames>("FListInputNames", \
150 [](const nnvm::NodeAttrs& attrs) { \
151 int num_args = nnvm::get<AtleastNDParam>(attrs.parsed).num_args; \
152 std::vector<std::string> ret; \
153 for (int i = 0; i < num_args; i++) { \
154 ret.push_back(std::string("ary") + std::to_string(i)); \
155 } \
156 return ret; \
157 }) \
158 .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<-1, -1>) \
159 .set_attr<mxnet::FInferShape>("FInferShape", AtleastNDShape<N>) \
160 .set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes) \
161 .set_attr<FCompute>("FCompute<cpu>", AtleastNDCompute<cpu>) \
162 .add_argument("arys", "NDArray-or-Symbol[]", "List of input arrays") \
163 .add_arguments(AtleastNDParam::__FIELDS__()) \
164
165 NNVM_REGISTER_ATLEAST_ND(1);
166
167 NNVM_REGISTER_ATLEAST_ND(2);
168
169 NNVM_REGISTER_ATLEAST_ND(3);
170
171 NNVM_REGISTER_OP(_npi_full_like)
172 .set_num_inputs(1)
173 .set_num_outputs(1)
174 .set_attr_parser(ParamParser<FullLikeOpParam>)
175 .set_attr<mxnet::FInferShape>("FInferShape", ElemwiseShape<1, 1>)
176 .set_attr<nnvm::FInferType>("FInferType", FullLikeOpType<FullLikeOpParam>)
177 .set_attr<nnvm::FIgnoreInputs>("FIgnoreInputs",
__anon883e59af0102(const NodeAttrs& attrs) 178 [](const NodeAttrs& attrs) {
179 return std::vector<uint32_t>(1, 0);
180 })
181 .set_attr<nnvm::FListInputNames>("FListInputNames",
__anon883e59af0202(const NodeAttrs& attrs) 182 [](const NodeAttrs& attrs) {
183 return std::vector<std::string>{"a"};
184 })
185 .set_attr<FCompute>("FCompute<cpu>", FullLikeOpCompute<cpu>)
186 .set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
187 .add_argument("a", "NDArray-or-Symbol",
188 "The shape and data-type of a define these same attributes of the returned array.")
189 .add_arguments(FullLikeOpParam::__FIELDS__());
190
191 NNVM_REGISTER_OP(_npi_arange)
192 .set_num_inputs(0)
193 .set_num_outputs(1)
194 .set_attr_parser(RangeParamParser)
195 .set_attr<mxnet::FInferShape>("FInferShape", NumpyRangeShape)
196 .set_attr<nnvm::FInferType>("FInferType", InitType<RangeParam>)
197 .set_attr<FCompute>("FCompute<cpu>", RangeCompute<cpu, RangeParam>)
198 .add_arguments(RangeParam::__FIELDS__());
199
200 NNVM_REGISTER_OP(_npi_eye)
201 .describe("Return a 2-D array with ones on the diagonal and zeros elsewhere.")
202 .set_num_inputs(0)
203 .set_num_outputs(1)
204 .set_attr_parser(ParamParser<NumpyEyeParam>)
205 .set_attr<mxnet::FInferShape>("FInferShape", NumpyEyeShape)
206 .set_attr<nnvm::FInferType>("FInferType", InitType<NumpyEyeParam>)
207 .set_attr<FCompute>("FCompute<cpu>", NumpyEyeFill<cpu>)
208 .add_arguments(NumpyEyeParam::__FIELDS__());
209
210 NNVM_REGISTER_OP(_npi_indices)
211 .describe("Return an array representing the indices of a grid.")
212 .set_num_inputs(0)
213 .set_num_outputs(1)
214 .set_attr_parser(ParamParser<IndicesOpParam>)
215 .set_attr<mxnet::FInferShape>("FInferShape", NumpyIndicesShape)
216 .set_attr<nnvm::FInferType>("FInferType", InitType<IndicesOpParam>)
217 .set_attr<FCompute>("FCompute<cpu>", IndicesCompute<cpu>)
218 .add_arguments(IndicesOpParam::__FIELDS__());
219
220 NNVM_REGISTER_OP(_npi_logspace)
221 .describe("Return numbers spaced evenly on a log scale.")
222 .set_num_inputs(0)
223 .set_num_outputs(1)
224 .set_attr_parser(ParamParser<LogspaceParam>)
225 .set_attr<mxnet::FInferShape>("FInferShape", LogspaceShape)
226 .set_attr<nnvm::FInferType>("FInferType", InitType<LogspaceParam>)
227 .set_attr<FCompute>("FCompute<cpu>", LogspaceCompute<cpu>)
228 .add_arguments(LogspaceParam::__FIELDS__());
229
230 } // namespace op
231 } // namespace mxnet
232