add large tensor test for linalg_gemm2

apache · Jul 15, 2020 · ab5cee5 · ab5cee5
1 parent e2366e9
commit ab5cee5
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py
@@ -1745,6 +1745,15 @@ def check_minimum():
     check_maximum()
     check_minimum()
 
+def test_linalg_ops():
+    def check_linalg_gemm2():
+        a= mx.nd.ones(shape=(SMALL_Y, LARGE_X))
+        b= mx.nd.ones(shape=(LARGE_X, SMALL_Y))
+        res=nd.linalg_gemm2(a, b)
+        res.shape == (SMALL_Y, SMALL_Y)
+        assert res.asnumpy()[0][0] == LARGE_X
+        assert res.asnumpy()[-1][-1] == LARGE_X
+    check_linalg_gemm2()
 
 @pytest.mark.timeout(0)
 def test_sparse_dot():