Zha0q1 commented on a change in pull request #18782: URL: https://github.com/apache/incubator-mxnet/pull/18782#discussion_r460171417
########## File path: tests/nightly/test_large_array.py ########## @@ -1791,6 +1792,29 @@ def test_sparse_dot(): assert out.shape == (2, 2) +def test_linalg_operators(): + def check_syrk_batch(): + # test both forward and backward + # batch syrk will be applied to the last two dimensions + A = nd.zeros((2, LARGE_SQ_X, LARGE_SQ_X)) + for i in range(LARGE_SQ_X): + A[0,i,i] = 1 + A[1,i,i] = 0.1 + A.attach_grad() + with mx.autograd.record(): + out = nd.linalg.syrk(A, alpha=2, transpose=False) + for i in range(LARGE_SQ_X): + assert out[0,i,i] == 2 + assert_almost_equal(out[1,i,i], nd.array([0.02]), rtol=1e-3, atol=1e-5) + out.backward() + for i in range(LARGE_SQ_X): + # check the first row + assert A.grad[0,0,i] == 4 + assert_almost_equal(A.grad[1,0,i], nd.array([0.4]), rtol=1e-3, atol=1e-5) Review comment: This is the correct result I believe. I verified with hand-written calculation. Yeah it also struck as counter-intuitive to me.. I am going to dive deep in matrix grad when I find time ---------------------------------------------------------------- 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