move matrix generator to utils

apache · Jul 23, 2020 · 8bbc53b · 8bbc53b
1 parent b615229
commit 8bbc53b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/python/mxnet/test_utils.py b/python/mxnet/test_utils.py
@@ -297,6 +297,13 @@ def create_vector(size, dtype=np.int64):
     a = mx.nd.arange(0, size, dtype=dtype)
     return a
 
+# For testing Large Square Matrix with total size > 2^32 elements
+def get_large_identity_mat():
+    A = nd.zeros((LARGE_SQ_X, LARGE_SQ_X))
+    for i in range(LARGE_SQ_X):
+        A[i,i] = 1
+    return A
+
 def rand_sparse_ndarray(shape, stype, density=None, dtype=None, distribution=None,
                         data_init=None, rsp_indices=None, modifier_func=None,
                         shuffle_csr_indices=False, ctx=None):

diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py
@@ -1170,12 +1170,6 @@ def check_correctness(mxnet_op, numpy_op, atol=1e-3):
 
 
 def test_linalg():
-    def get_large_identity_mat():
-        A = nd.zeros((LARGE_SQ_X, LARGE_SQ_X))
-        for i in range(LARGE_SQ_X):
-            A[i,i] = 1
-        return A
-
     def batchify(mat):
         return nd.array([mat.asnumpy()])