Add syrk test shape check (#18812)

* add shape check * add name to contributor.md Co-authored-by: Ubuntu <ubuntu@ip-172-31-6-47.us-west-2.compute.internal>
apache · Jul 29, 2020 · 85eb528 · 85eb528
1 parent 7bef9cb
commit 85eb528
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git 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
@@ -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)