diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py index 9dc29eb1584e..03ecdb88f84f 100644 --- a/tests/nightly/test_large_array.py +++ b/tests/nightly/test_large_array.py @@ -291,10 +291,10 @@ def test_unravel_index(): assert (indices_2d.asnumpy() == np.array(original_2d_indices)).all() -def create_2d_tensor(rows, columns): +def create_2d_tensor(rows, columns, dtype=np.int64): a = np.arange(0, rows).reshape(rows, 1) b = np.broadcast_to(a, shape=(a.shape[0], columns)) - return nd.array(b, dtype=np.int64) + return nd.array(b, dtype=dtype) def test_transpose(): @@ -351,6 +351,113 @@ def test_topk(): assert l.sum() == np.sum(np.arange(0, SMALL_Y)) +def test_add(): + a = nd.ones(shape=(LARGE_X, SMALL_Y)) + b = nd.ones(shape=(LARGE_X, SMALL_Y)) + c = b + c = c.__add__(a) + assert c[0][-1] == 2 + assert c.shape == a.shape + + +def test_sub(): + a = 3*nd.ones(shape=(LARGE_X, SMALL_Y)) + b = nd.ones(shape=(LARGE_X, SMALL_Y)) + c = b + c = c.__sub__(a) + assert c[0][-1] == -2 + assert c.shape == a.shape + + +def test_rsub(): + a = 3*nd.ones(shape=(LARGE_X, SMALL_Y)) + b = nd.ones(shape=(LARGE_X, SMALL_Y)) + c = b + c = c.__rsub__(a) + assert c[0][-1] == 2 + assert c.shape == a.shape + + +def test_neg(): + a = nd.ones(shape=(LARGE_X, SMALL_Y)) + c = a + c = c.__neg__() + assert c[0][-1] == -1 + assert c.shape == a.shape + + +def test_mul(): + a = 2*nd.ones(shape=(LARGE_X, SMALL_Y)) + b = 3*nd.ones(shape=(LARGE_X, SMALL_Y)) + c = b + c = c.__mul__(a) + assert c[0][-1] == 6 + assert c.shape == a.shape + + +def test_div(): + a = 2*nd.ones(shape=(LARGE_X, SMALL_Y)) + b = 3*nd.ones(shape=(LARGE_X, SMALL_Y)) + c = b + c = c.__div__(a) + assert c[0][-1] == 3/2 + assert c.shape == a.shape + + +def test_rdiv(): + a = 2*nd.ones(shape=(LARGE_X, SMALL_Y)) + b = 3*nd.ones(shape=(LARGE_X, SMALL_Y)) + c = b + c = c.__rdiv__(a) + assert c[0][-1] == 2/3 + assert c.shape == a.shape + + +def test_mod(): + a = 2*nd.ones(shape=(LARGE_X, SMALL_Y)) + b = 3*nd.ones(shape=(LARGE_X, SMALL_Y)) + c = b + c = c.__mod__(a) + assert c[0][-1] == 1 + assert c.shape == a.shape + + +def test_rmod(): + a = 2*nd.ones(shape=(LARGE_X, SMALL_Y)) + b = 3*nd.ones(shape=(LARGE_X, SMALL_Y)) + c = b + c = c.__rmod__(a) + assert c[0][-1] == 2 + assert c.shape == a.shape + + +def test_imod(): + a = 2*nd.ones(shape=(LARGE_X, SMALL_Y)) + b = 3*nd.ones(shape=(LARGE_X, SMALL_Y)) + c = b + c = c.__imod__(a) + assert c[0][-1] == 1 + assert c.shape == a.shape + + +def test_pow(): + a = 2*nd.ones(shape=(LARGE_X, SMALL_Y)) + b = 3*nd.ones(shape=(LARGE_X, SMALL_Y)) + c = b + c = c.__pow__(a) + assert c[0][-1] == 9 + assert c.shape == a.shape + + +def test_rpow(): + a = 2*nd.ones(shape=(LARGE_X, SMALL_Y)) + b = 3*nd.ones(shape=(LARGE_X, SMALL_Y)) + c = b + c = c.__rpow__(a) + assert c[0][-1] == 8 + assert c.shape == a.shape + + if __name__ == '__main__': import nose nose.runmodule()