Add Large Tensor Test for linalg_syrk (#18782)

* add large tensor test for syrk, foward and backward

* change to batch input

* move syrk test into test-linalg

Co-authored-by: Ubuntu <[email protected]>

tests/nightly/test_large_array.py

@@ -37,7 +37,7 @@
 LARGE_X = 100000000
 SMALL_X = 100
 SMALL_Y = 50
-LARGE_SQ_X = 80000
+LARGE_SQ_X = 70000
 LARGE_SIZE = LARGE_X * SMALL_Y
 LARGE_TENSOR_SHAPE = 2**32
 RNN_LARGE_TENSOR = 2**28
@@ -1190,9 +1190,26 @@ def check_potri():
         # output should be an identity matrix
         for i in range(LARGE_SQ_X):
             assert out[i,i] == 1
+
+    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)
+        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[0,0,0] == 4
+        assert_almost_equal(A.grad[1,0,0], nd.array([0.4]), rtol=1e-3, atol=1e-5)
 
     check_potrf()
     check_potri()
+    check_syrk_batch()
 
 
 def test_basic():