D-Roberts commented on a change in pull request #18197:
URL: https://github.com/apache/incubator-mxnet/pull/18197#discussion_r425140133
##########
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:
@hzfan Thank you for your review. Added comments and updated reference
in the code.
----------------------------------------------------------------
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
