hzfan commented on a change in pull request #18197:
URL: https://github.com/apache/incubator-mxnet/pull/18197#discussion_r424862190
##########
File path: src/operator/numpy/linalg/np_qr-inl.h
##########
@@ -514,15 +548,28 @@ struct qr_backward {
template<typename xpu>
size_t QrBackwardWorkspaceSize(const TBlob& a,
+ const TBlob& q,
const TBlob& r,
const TBlob& grad_a) {
+ const mxnet::TShape& a_shape = a.shape_;
+ const int a_ndim = a_shape.ndim();
+ const int n = a.size(a_ndim - 1);
+ const int m = a.size(a_ndim - 2);
+
if (0U == a.Size()) { return 0U; }
MSHADOW_SGL_DBL_TYPE_SWITCH(grad_a.type_flag_, DType, {
size_t work_space_size = 0;
- // for grad a and M
work_space_size += a.Size();
- work_space_size += r.Size();
+ if (m >= n) {
+ work_space_size += r.Size();
+ } else {
+ const mxnet::TShape& q_shape = q.shape_;
+ mxnet::TShape v_shape(q_shape);
+ v_shape[a_ndim - 1] = n - m;
+ work_space_size += 5 * q.Size();
Review comment:
Could you add some documents about what does `5` and `3` here stand for?
Seem like magic numbers.
----------------------------------------------------------------
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:
[email protected]
