Large tensor support for random ops

tests/nightly/test_large_array.py

@@ -24,6 +24,7 @@
 # dimension constants
 MEDIUM_X = 10000
 LARGE_X = 100000000
+SMALL_X = 100
 SMALL_Y = 50
 LARGE_SIZE = LARGE_X * SMALL_Y
 
@@ -79,6 +80,88 @@ def test_ndarray_random_randint():
     assert a.__gt__(low) and a.__lt__(high)
 
 
+@with_seed()
+def test_ndarray_random_exponential():
+    scale_array = nd.random.uniform(shape=(MEDIUM_X, SMALL_Y))
+    a = nd.random.exponential(scale=scale_array, shape=(SMALL_X, SMALL_Y))
+    assert a.shape == (MEDIUM_X, SMALL_Y, SMALL_X, SMALL_Y)
+    assert a[-1][0][0][0] >= 0
+
+
+@with_seed()
+def test_ndarray_random_gamma():
+    alpha_array = nd.random.uniform(shape=(MEDIUM_X, SMALL_Y))
+    beta_array = nd.random.uniform(shape=(MEDIUM_X, SMALL_Y))
+    a = nd.random.gamma(alpha=alpha_array, beta=beta_array,
+                        shape=(SMALL_X, SMALL_Y))
+    assert a.shape == (MEDIUM_X, SMALL_Y, SMALL_X, SMALL_Y)
+    assert a[-1][0][0][0] >= 0
+
+
+@with_seed()
+def test_ndarray_random_multinomial():
+    # test 1 shape dimension
+    probs = nd.random.uniform(shape=(LARGE_X, SMALL_Y))
+    a = nd.random.multinomial(probs)
+    assert a.shape == (LARGE_X,)
+    assert a[-1] >= 0
+    # test for NDArray multi-dimension shape
+    a = nd.random.multinomial(probs, shape=(SMALL_X, SMALL_Y))
+    assert a.shape == (LARGE_X, SMALL_X, SMALL_Y)
+    assert a[-1][0][0] >= 0
+    # test log_likelihood output shape
+    a = nd.random.multinomial(probs, shape=(SMALL_X, SMALL_Y), get_prob=True)
+    assert a[0].shape == (LARGE_X, SMALL_X, SMALL_Y) and a[0].shape == a[1].shape
+    assert a[-1][0][0] >= 0
+
+
+@with_seed()
+def test_ndarray_random_generalized_negative_binomial():
+    alpha_array = nd.random.uniform(shape=(MEDIUM_X, SMALL_Y))
+    mu_array = nd.random.uniform(shape=(MEDIUM_X, SMALL_Y))
+    a = nd.random.generalized_negative_binomial(mu=mu_array, alpha=alpha_array,
+                                                shape=(SMALL_X, SMALL_Y))
+    assert a.shape == (MEDIUM_X, SMALL_Y, SMALL_X, SMALL_Y)
+    assert a[-1][0][0][0] >= 0
+
+
+@with_seed()
+def test_ndarray_random_negative_binomial():
+    k_array = nd.random.uniform(shape=(MEDIUM_X, SMALL_Y))
+    p_array = nd.random.uniform(shape=(MEDIUM_X, SMALL_Y))
+    a = nd.random.negative_binomial(k=k_array, p=p_array,
+                                    shape=(SMALL_X, SMALL_Y))
+    assert a.shape == (MEDIUM_X, SMALL_Y, SMALL_X, SMALL_Y)
+    assert a[-1][0][0][0] >= 0
+
+
+@with_seed()
+def test_ndarray_random_normal():
+    scale_array = nd.random.uniform(shape=(MEDIUM_X, SMALL_Y))
+    loc_array = nd.random.uniform(shape=(MEDIUM_X, SMALL_Y))
+    a = nd.random.normal(loc=loc_array, scale=scale_array,
+                         shape=(SMALL_X, SMALL_Y))
+    assert a.shape == (MEDIUM_X, SMALL_Y, SMALL_X, SMALL_Y)
+    assert a[-1][0][0][0] >= 0
+
+
+@with_seed()
+def test_ndarray_random_poisson():
+    lambda_array = nd.random.uniform(shape=(MEDIUM_X, SMALL_Y))
+    a = nd.random.poisson(lam=lambda_array, shape=(SMALL_X, SMALL_Y))
+    assert a.shape == (MEDIUM_X, SMALL_Y, SMALL_X, SMALL_Y)
+    assert a[-1][0][0][0] >= 0
+
+
+@with_seed()
+def test_ndarray_random_randn():
+    a = nd.random.randn(LARGE_X, SMALL_Y)
+    assert a.shape == (LARGE_X, SMALL_Y)
+    assert a[-1][0] >= 0
+    # TODO: Once PR for randn ndarray dtype for loc,scale param merged
+    # Add check for (x,y,m,n) where x,y shape of loc,scale and m,n input shape
+
+
 def test_ndarray_empty():
     a = nd.empty((LARGE_X, SMALL_Y))
     assert a.shape == (LARGE_X, SMALL_Y)