diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index bd7f966aaa5f..be04d82ffe92 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -252,6 +252,7 @@ List of Contributors * [Oliver Kowalke](https://github.com/olk) * [Connor Goggins](https://github.com/connorgoggins) * [Joe Evans](https://github.com/josephevans) +* [Zhaoqi Zhu](https://github.com/zha0q1) Label Bot --------- diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py index 306c827bab9f..8865eae2b81e 100644 --- a/tests/nightly/test_large_array.py +++ b/tests/nightly/test_large_array.py @@ -1201,9 +1201,11 @@ def check_syrk_batch(): A.attach_grad() with mx.autograd.record(): out = nd.linalg.syrk(A, alpha=2, transpose=False) + assert out.shape == (2, LARGE_SQ_X, LARGE_SQ_X) assert out[0,0,0] == 2 assert_almost_equal(out[1,0,0], nd.array([0.02]), rtol=1e-3, atol=1e-5) out.backward() + assert A.grad.shape == (2, LARGE_SQ_X, LARGE_SQ_X) assert A.grad[0,0,0] == 4 assert_almost_equal(A.grad[1,0,0], nd.array([0.4]), rtol=1e-3, atol=1e-5)