add shape checks

apache · Jul 27, 2020 · e170000 · e170000
1 parent 629175f
commit e170000
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py
@@ -1220,8 +1220,10 @@ def run_gemm2(inp1,inp2):
         inp2=mx.nd.ones(shape=(LARGE_X, SMALL_Y))
         inp1_grad, inp2_grad, out= run_gemm2(inp1,inp2)
         assert out.asnumpy()[0][0] == LARGE_X
+        assert out.shape() == (SMALL_Y, SMALL_Y)
         out.backward()
-        assert inp1_grad.asnumpy()[0][0] == SMALL_Y
+        assert inp1_grad.asnumpy()[0][0] == (SMALL_Y, LARGE_X)
+        assert inp2_grad.asnumpy()[0][0] == (LARGE_X, SMALL_Y)
         assert_almost_equal(inp2_grad.asnumpy()[0][0],49.1)