haojin2 commented on a change in pull request #16720: [Numpy] Implement numpy operator 'average' URL: https://github.com/apache/incubator-mxnet/pull/16720#discussion_r341960103
########## File path: src/operator/numpy/np_broadcast_reduce_op.h ########## @@ -398,6 +399,353 @@ void ReduceAxesComputeWithWorkspaceImpl(const OpContext& ctx, }); } +struct NumpyWeightedAverageParam : public dmlc::Parameter<NumpyWeightedAverageParam> { + dmlc::optional<mxnet::Tuple<int>> axis; + bool returned; + bool weighted; + + DMLC_DECLARE_PARAMETER(NumpyWeightedAverageParam) { + DMLC_DECLARE_FIELD(axis) + .set_default(dmlc::optional<mxnet::Tuple<int>>()) + .describe("Axis or axes along which a average is performed. The default, axis=None, will average " + "all of the elements of the input array. If axis is negative it counts from the " + "last to the first axis."); + DMLC_DECLARE_FIELD(returned) + .set_default(false) + .describe("If True, the tuple (average, sum_of_weights) is returned," + "otherwise only the average is returned." + "If weights=None, sum_of_weights is equivalent to" + "the number of elements over which the average is taken."); + DMLC_DECLARE_FIELD(weighted) + .set_default(true) + .describe("Auxiliary flag to deal with none weights."); + } +}; + +inline bool NumpyWeightedAverageShape(const nnvm::NodeAttrs& attrs, + std::vector<TShape> *in_attrs, + std::vector<TShape> *out_attrs) { + const NumpyWeightedAverageParam& param = nnvm::get<NumpyWeightedAverageParam>(attrs.parsed); + CHECK_EQ(in_attrs->size(), (param.weighted ? 2U : 1U)); + CHECK_EQ(out_attrs->size(), 2U); + if (!shape_is_known(in_attrs->at(0))) { + return false; + } + + const TShape& a_shape = (*in_attrs)[0]; + SHAPE_ASSIGN_CHECK(*out_attrs, 0, + NumpyReduceAxesShapeImpl(a_shape, param.axis, false)); + + if (param.weighted) { + const TShape& w_shape = (*in_attrs)[1]; + if (w_shape.ndim() != a_shape.ndim()) { + CHECK_EQ(w_shape.ndim(), 1U) << "1D weights expected when shapes of a and weights differ."; + + CHECK_EQ(param.axis.has_value(), true) << "Axis must be specified when shapes of a and weights differ."; + + mxnet::Tuple<int> axes(param.axis.value()); + + CHECK_EQ(axes.ndim(), 1U) << "Axis must be int when shapes of a and weights differ."; + + int red_axis = axes[0] < 0 ? axes[0] + a_shape.ndim() : axes[0]; + + CHECK_EQ(a_shape[red_axis], w_shape[0]) << "Length of weights not compatible with specified " + "axis."; + + SHAPE_ASSIGN_CHECK(*out_attrs, 1, + NumpyReduceAxesShapeImpl(w_shape, dmlc::optional<mxnet::Tuple<int>>(), false)); + } else { + for (int i = 0; i < w_shape.ndim(); i++) { + CHECK_EQ(w_shape[i], a_shape[i]); + } + SHAPE_ASSIGN_CHECK(*out_attrs, 1, + NumpyReduceAxesShapeImpl(w_shape, param.axis, false)); + } + } else { + SHAPE_ASSIGN_CHECK(*out_attrs, 1, TShape(0, -1)); + } + + return shape_is_known(out_attrs->at(0)) && shape_is_known(out_attrs->at(1)); +} + +template<int req, bool onedim = false> +struct avg_grad_a_kernel { + template<typename DType> + MSHADOW_XINLINE static void Map(int i, + DType* out, + const DType* w, + const DType* scl, + const DType* ograd, + const mshadow::Shape<6>& small, + const mshadow::Shape<6>& big) { + // partial a = w / sum(w) + size_t big_idx = i; + size_t small_idx = i; + size_t big_stride = 1; + size_t small_stride = 1; + size_t red_axis_idx = 0; + for (int axis = 5; axis >= 0; --axis) { + size_t axis_idx = big_idx % big[axis]; + small_idx -= axis_idx * big_stride; + if (small[axis] != 1) { + small_idx += axis_idx * small_stride; + } else if (onedim && small[axis] != big[axis]) { + red_axis_idx = axis_idx; + } + big_idx /= big[axis]; + big_stride *= big[axis]; + small_stride *= small[axis]; + } + if (onedim) { + KERNEL_ASSIGN(out[i], req, (ograd[small_idx] * (w[red_axis_idx] / *scl))); + } else { + KERNEL_ASSIGN(out[i], req, (ograd[small_idx] * (w[i] / scl[small_idx]))); + } + } +}; + +template<int req> +struct avg_grad_w_kernel { + template<typename DType> + MSHADOW_XINLINE static void Map(int i, + DType* out, + const DType* a, + const DType* scl, + const DType* sum_of_wa, + const DType* ograd, + const mshadow::Shape<6>& small, + const mshadow::Shape<6>& big) { + // partial w = (a * sum(w) - sum(a*w)) / (sum(w) * sum(w)) + size_t big_idx = i; + size_t small_idx = i; + size_t big_stride = 1; + size_t small_stride = 1; + for (int axis = 5; axis >= 0; --axis) { + size_t axis_idx = big_idx % big[axis]; + small_idx -= axis_idx * big_stride; + if (small[axis] != 1) { + small_idx += axis_idx * small_stride; + } + big_idx /= big[axis]; + big_stride *= big[axis]; + small_stride *= small[axis]; + } + DType ret = ograd[small_idx] * (((a[i] * scl[small_idx] - sum_of_wa[small_idx]) / scl[small_idx]) / scl[small_idx]); + KERNEL_ASSIGN(out[i], req, ret); + } +}; + +template<int req> +struct avg_grad_w_1D_kernel { + template<typename DType> + MSHADOW_XINLINE static void Map(int i, + DType* out, + const DType* a, + const DType* scl, + const DType* sum_of_wa, + const DType* ograd, + const mshadow::Shape<6>& big, + const int red_axis) { + DType scl_val = *scl; + size_t tail = 1; + size_t head = 1; + for (int axis = 5; axis > red_axis; --axis) { + tail *= big[axis]; + } + for (int axis = 0; axis < red_axis; ++axis) { + head *= big[axis]; + } + DType ret = 0; + for (size_t j = 0; j < head; ++j) { + for (size_t k = 0; k < tail; ++k) { + size_t a_idx = j*(tail*big[red_axis]) + i * tail + k; + size_t small_idx = j*tail + k; + ret += (ograd[small_idx] * (((a[a_idx] * scl_val - sum_of_wa[small_idx]) / scl_val) / scl_val)); + } + } + KERNEL_ASSIGN(out[i], req, ret); + } +}; + +inline mshadow::Shape<6> TShapeToMShadowShape(const mxnet::TShape& shape) { + mshadow::Shape<6> mshape; + for (int i = 0; i < 6; ++i) { + mshape[i] = (i < shape.ndim()) ? shape[i] : 1; + } + return mshape; +} + +template<typename xpu, bool back = false> +void NumpyWeightedAverageComputeImpl(const nnvm::NodeAttrs& attrs, + const OpContext& ctx, + const std::vector<TBlob>& inputs, + const std::vector<OpReqType>& req, + const std::vector<TBlob>& outputs, + const dmlc::optional<mxnet::Tuple<int>>& axis) { + using namespace mshadow; + using namespace mxnet_op; + Stream<xpu>* s = ctx.get_stream<xpu>(); + const TBlob& data = inputs[0]; + TShape small1 = NumpyReduceAxesShapeImpl(data.shape_, axis, true); + // Reshape weights + TShape small2 = small1; + TBlob weights = inputs[1]; + + bool one_dim = weights.shape_.ndim() != data.shape_.ndim(); + + int red_axis = -1; + + if (one_dim) { + CHECK_EQ(weights.shape_.ndim(), 1U) << "1D weights expected when shapes of a and weights differ."; + Review comment: Same here, no need for extra blank lines in between every pair of statements. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services