haojin2 commented on a change in pull request #16786: Add OP diag [numpy]
URL: https://github.com/apache/incubator-mxnet/pull/16786#discussion_r347012799
##########
File path: src/operator/numpy/np_matrix_op-inl.h
##########
@@ -945,6 +950,206 @@ void NumpyConcatenateBackward(const nnvm::NodeAttrs&
attrs,
});
}
+struct NumpyDiagParam : public dmlc::Parameter<NumpyDiagParam> {
+ int k;
+ DMLC_DECLARE_PARAMETER(NumpyDiagParam) {
+ DMLC_DECLARE_FIELD(k).set_default(0)
+ .describe("Diagonal in question. The default is 0. "
+ "Use k>0 for diagonals above the main diagonal, "
+ "and k<0 for diagonals below the main diagonal. ");
+ }
+};
+
+inline mxnet::TShape NumpyDiagShapeImpl(const mxnet::TShape &ishape,
+ const int k) {
+ CHECK_LE(ishape.ndim(), 2) << "Input must be 1- or 2-d";
+
+ if (ishape.ndim() == 1) {
+ auto s = ishape[0] + std::abs(k);
+ return mxnet::TShape({s, s});
+ }
+
+ auto h = ishape[0];
+ auto w = ishape[1];
+
+ if (k > 0) {
+ w -= k;
+ } else if (k < 0) {
+ h += k;
+ }
+ dim_t a = 0;
+ auto s = std::max(std::min(h, w), a);
+ // s is the length of diagonal with k as the offset
+
+ int32_t n_dim = ishape.ndim() - 1;
+ mxnet::TShape oshape(n_dim, -1);
+ oshape[n_dim - 1] = s;
+ return oshape;
+}
+
+inline bool NumpyDiagOpShape(const nnvm::NodeAttrs &attrs,
+ mxnet::ShapeVector *in_attrs,
+ mxnet::ShapeVector *out_attrs) {
+ CHECK_EQ(in_attrs->size(), 1U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ const mxnet::TShape &ishape = (*in_attrs)[0];
+ if (!mxnet::ndim_is_known(ishape)) {
+ return false;
+ }
+
+ const NumpyDiagParam ¶m = nnvm::get<NumpyDiagParam>(attrs.parsed);
+ mxnet::TShape oshape = NumpyDiagShapeImpl(ishape, param.k);
+
+ if (shape_is_none(oshape)) {
+ LOG(FATAL) << "Diagonal does not exist.";
+ }
+ SHAPE_ASSIGN_CHECK(*out_attrs, 0, oshape);
+
+ return shape_is_known(out_attrs->at(0));
+}
+
+inline bool NumpyDiagOpType(const nnvm::NodeAttrs &attrs,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ CHECK_EQ(in_attrs->size(), 1U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ TYPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[0]);
+ TYPE_ASSIGN_CHECK(*in_attrs, 0, (*out_attrs)[0]);
+ return (*out_attrs)[0] != -1;
+}
+
+template <int ndim, int req, bool back>
+struct diag {
+ template <typename DType>
+ MSHADOW_XINLINE static void Map(index_t i, DType *out, const DType *a,
+ index_t stride, index_t offset) {
+ using namespace mxnet_op;
+ index_t j = offset + stride * i;
+
+ if (back) {
+ KERNEL_ASSIGN(out[j], req, a[i]);
+ } else {
+ KERNEL_ASSIGN(out[i], req, a[j]);
+ }
+ }
+};
+
+template <int req, bool back>
+struct diag_gen {
+ template <typename DType>
+ MSHADOW_XINLINE static void Map(index_t i, DType *out, const DType *a,
+ mshadow::Shape<2> oshape, int k) {
+ using namespace mxnet_op;
+
+ auto j = unravel(i, oshape);
+ if (j[1] == (j[0] + k)) {
+ auto l = j[0] < j[1] ? j[0] : j[1];
+ if (back) {
+ KERNEL_ASSIGN(out[l], req, a[i]);
+ } else {
+ KERNEL_ASSIGN(out[i], req, a[l]);
+ }
+ } else if (!back) {
+ KERNEL_ASSIGN(out[i], req, static_cast<DType>(0));
+ }
+ }
+};
+
+template <typename xpu, bool back>
+void NumpyDiagOpImpl(const TBlob &in_data,
+ const TBlob &out_data,
+ const mxnet::TShape &ishape,
+ const mxnet::TShape &oshape,
+ index_t dsize,
+ const int &k,
+ mxnet_op::Stream<xpu> *s,
+ const std::vector<OpReqType> &req) {
Review comment:
As I suggested last time, simply pass the `OpReqType` value instead of the
whole vector.
----------------------------------------------------------------
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
