test_large_vector.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import os
import sys
import tempfile
import math
import numpy as np
import mxnet as mx

curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
sys.path.append(os.path.join(curr_path, '../python/unittest/'))

from mxnet.test_utils import rand_ndarray, assert_almost_equal, rand_coord_2d, create_vector
from mxnet import gluon, nd
from common import with_seed
import pytest


# dimension constants
LARGE_X = 4300000000
MEDIUM_X = 1000000000


@pytest.mark.timeout(0)
def test_nn():
    def check_dense():
        data = mx.nd.ones(shape=LARGE_X)
        linear = gluon.nn.Dense(2)
        linear.initialize()
        res = linear(data)
        assert res.shape == (LARGE_X, 2)

    def check_sign():
        a = mx.nd.random.normal(-1, 1, shape=LARGE_X)
        mx_res = mx.nd.sign(a)
        assert_almost_equal(mx_res[-1].asnumpy(), np.sign(a[-1].asnumpy()))

    # TODO: correctness of layernorm
    # numpy implementation for large vector is flaky
    def check_layer_norm():
        axis = 0
        eps = 1E-5
        in_shape = LARGE_X
        data = nd.random.normal(0, 1, in_shape)
        gamma = nd.random.normal(0, 1, in_shape)
        beta = nd.random.normal(0, 1, in_shape)
        mx_out = nd.LayerNorm(data, gamma, beta, axis, eps)
        assert mx_out.shape == (in_shape,)

    # TODO: correctness of batchnorm
    # in future, we could test if mean, var of output
    # matches target output's mean, var
    def check_batchnorm():
        shape = LARGE_X
        axis = 0  # since vector
        data = mx.nd.ones(shape=shape)
        bn_gamma = mx.nd.random.uniform(shape=shape)
        bn_beta = mx.nd.random.uniform(shape=shape)
        bn_running_mean = mx.nd.zeros(shape)
        bn_running_var = mx.nd.ones(shape)
        output = mx.nd.BatchNorm(data, bn_gamma, bn_beta,
                                 bn_running_mean, bn_running_var, axis=axis)
        assert output.shape == (shape,)

    def check_sequence_mask():
        # Sequence Mask input [max_sequence_length, batch_size]
        # test with input batch_size = 2
        a = nd.arange(0, LARGE_X * 2).reshape(LARGE_X, 2)
        # test as identity operator
        b = nd.SequenceMask(a)
        assert b[-1][0] == a[-1][0]
        assert b.shape == a.shape
        # test with default mask
        b = nd.SequenceMask(a, sequence_length=nd.array([1, 1]),
                            use_sequence_length=True)
        assert b[0][1] == a[0][1]  # first sequence of each batch kept
        assert b[-1][-1] != a[-1][-1]  # rest sequences masked
        assert b[-1][-1] == 0
        # test with mask value
        b = nd.SequenceMask(a, sequence_length=nd.array([1, 1]),
                            use_sequence_length=True, value=-1)
        assert b[-1][-1] == -1

    def check_sequence_reverse():
        a = nd.arange(0, LARGE_X * 2).reshape(LARGE_X, 2)
        # test as reverse operator
        b = nd.SequenceReverse(a)
        assert b[-1][0] == a[0][0]
        assert b.shape == a.shape
        # test with sequence length
        b = nd.SequenceReverse(a, sequence_length=nd.array([2, 3]),
                               use_sequence_length=True)
        assert b[1][0] == a[0][0]  # check if reversed
        assert b[-1][0] == a[-1][0]  # check if intact
        assert b.shape == a.shape

    def check_sequence_last():
        a = nd.arange(0, LARGE_X * 2).reshape(LARGE_X, 2)
        # test if returns last sequence
        b = nd.SequenceLast(a)
        assert_almost_equal(b.asnumpy(), a[-1].asnumpy())
        assert b.shape == (2,)
        # test with sequence length
        # parameter sequence_length - NDArray with shape (batch_size)
        # (2,3) indicates 2nd sequence from batch 1 and 3rd sequence from batch 2
        # need to mention dtype = int64 for sequence_length ndarray to support large indices
        # else it defaults to float32 and errors
        b = nd.SequenceLast(a, sequence_length=mx.nd.array([2, 3], dtype="int64"),
                            use_sequence_length=True)
        # check if it takes 2nd sequence from the first batch
        assert b[0] == a[1][0]

    check_sequence_last()
    check_dense()
    check_sign()
    check_layer_norm()
    check_batchnorm()
    check_sequence_mask()
    check_sequence_reverse()


@pytest.mark.timeout(0)
def test_tensor():
    def check_ndarray_zeros():
        a = nd.zeros(shape=LARGE_X)
        assert a[-1] == 0
        assert a.shape == (LARGE_X,)
        assert a.size == LARGE_X

    def check_ndarray_ones():
        a = nd.ones(shape=LARGE_X)
        assert a[-1] == 1
        assert nd.sum(a) == LARGE_X

    def check_ndarray_empty():
        a = nd.empty(LARGE_X)
        assert a.shape == (LARGE_X,)

    @with_seed()
    def check_ndarray_random_uniform():
        a = nd.random.uniform(shape=LARGE_X)
        assert a[-1] != 0

    @pytest.mark.skip(reason="Randint flaky, tracked at https://github.com/apache/incubator-mxnet/issues/16172")
    @with_seed()
    def check_ndarray_random_randint():
        # check if randint can generate value greater than 2**32 (large)
        low = 2**32
        high = 2**34
        a = nd.random.randint(low, high, dtype=np.int64, shape=LARGE_X).asnumpy()
        assert a.shape == (LARGE_X,)
        assert (a >= low).all() and (a < high).all()

    @with_seed()
    def check_ndarray_random_exponential():
        a = nd.random.exponential(shape=LARGE_X)
        assert a[-1] >= 0.
        assert a.shape[0] == LARGE_X

    @with_seed()
    def check_ndarray_random_gamma():
        a = nd.random.gamma(shape=LARGE_X)
        assert a[-1] >= 0.
        assert a.shape[0] == LARGE_X

    @with_seed()
    def check_ndarray_random_generalized_negative_binomial():
        a = nd.random.generalized_negative_binomial(shape=LARGE_X)
        assert a[-1] >= 0.
        assert a.shape[0] == LARGE_X

    @with_seed()
    def check_ndarray_random_multinomial():
        a = nd.random.multinomial(nd.random.uniform(shape=LARGE_X))
        assert a[-1] >= 0.
        assert a.shape[0] == 1

    @with_seed()
    def check_ndarray_random_negative_binomial():
        a = nd.random.negative_binomial(shape=LARGE_X)
        assert a[-1] >= 0.
        assert a.shape[0] == LARGE_X

    @with_seed()
    def check_ndarray_random_normal():
        a = nd.random.normal(shape=LARGE_X)
        assert a.shape[0] == LARGE_X

    @with_seed()
    def check_ndarray_random_poisson():
        a = nd.random.poisson(shape=LARGE_X)
        assert a[-1] >= 0.
        assert a.shape[0] == LARGE_X

    @with_seed()
    def check_ndarray_random_randn():
        a = nd.random.randn(LARGE_X)
        assert a.shape[0] == LARGE_X

    @with_seed()
    def check_ndarray_random_shuffle():
        a = nd.ones(shape=LARGE_X)
        a[-1] = 3
        a = nd.random.shuffle(a)
        unique_a = np.unique(a.asnumpy())
        assert len(unique_a) == 2  # only 2 unique values
        assert unique_a[0] == 1  # first unique value is 1
        assert unique_a[1] == 3  # second unique value is 3
        assert a.shape[0] == LARGE_X

    def check_full():
        a = nd.full(LARGE_X, 3)
        assert a.shape[0] == LARGE_X
        assert a[LARGE_X // 2] == 3
        assert a[-1] == 3

    def check_repeat():
        x = create_vector(size=LARGE_X//2)
        y = nd.repeat(x, repeats=2, axis = 0)
        assert y.shape[0] == LARGE_X
        assert y[1] == 0
        assert y[LARGE_X-1] == LARGE_X//2-1

    def check_clip():
        a = create_vector(LARGE_X)
        res = nd.clip(a, a_min=100, a_max=1000)
        assert res[-1] == 1000

    def check_slice():
        a = nd.ones(LARGE_X)
        res = nd.slice(a, begin=(LARGE_X - MEDIUM_X), end=LARGE_X)
        assert res.shape[0] == MEDIUM_X
        assert res[0] == 1

    def check_slice_assign():
        a = nd.ones(shape=LARGE_X)
        a[LARGE_X-1:LARGE_X] = 1000
        assert np.sum(a[-1].asnumpy() == 1000) == 1

    def check_take():
        a = nd.ones(shape=LARGE_X)
        idx = nd.arange(LARGE_X - 1000, LARGE_X)
        res = nd.take(a, idx)
        assert np.sum(res.asnumpy() == 1) == res.shape[0]

    def check_expand_dims():
        a = nd.ones(shape=LARGE_X)
        res = nd.expand_dims(a, axis=0)
        assert res[0][0] == 1
        assert res.shape == (1, a.shape[0])

    def check_squeeze():
        a = nd.ones(shape=LARGE_X)
        data = nd.expand_dims(a, axis=0)
        res = nd.squeeze(data)
        assert a[0] == res[0]
        assert res.shape == a.shape

    def check_broadcast_div():
        a = nd.ones(shape=LARGE_X)
        b = nd.ones(shape=LARGE_X) * 2
        res = a / b
        assert np.sum(res.asnumpy() == 0.5) == a.shape[0]

    def check_size():
        b = create_vector(size=LARGE_X)
        # explicit wait_to_read()
        assert b[0] == 0
        assert b.size == LARGE_X

    def check_copy():
        a = nd.ones(LARGE_X)
        b = a.copy()
        assert a[0] == b[0]
        assert b.shape == a.shape
        assert b.size == LARGE_X

    def check_copy_to():
        a = create_vector(size=LARGE_X)
        # keeping dtype same as input uses parallel copy which is much faster
        b = nd.zeros(LARGE_X, dtype=np.int64)
        c = a.copyto(b)
        assert c is b
        assert b[-1] == LARGE_X-1
        assert b[0] == 0

    def check_zeros_like():
        a = nd.ones(LARGE_X)
        b = nd.zeros_like(a)
        assert b[-1] == 0
        assert b.shape == a.shape

    def check_ones_like():
        a = nd.zeros(LARGE_X)
        b = nd.ones_like(a)
        assert b[-1] == 1
        assert b.shape == a.shape

    def check_shape():
        b = create_vector(size=LARGE_X)
        # explicit wait_to_read()
        assert b[0] == 0
        assert b.shape[0] == LARGE_X

    def check_concat():
        a = nd.ones(LARGE_X)
        b = nd.zeros(LARGE_X)
        c = nd.concat(a, b, dim=0)
        assert c[0] == 1
        assert c[-1] == 0
        assert c.shape[0] == (2 * LARGE_X)

    def check_slice_like():
        a = create_vector(size=LARGE_X)
        b = nd.ones(LARGE_X//2)
        c = nd.slice_like(a, b)
        assert c.shape == b.shape
        assert c[0] == 0
        assert c[-1] == (LARGE_X // 2 - 1)

    def check_slice_axis():
        a = create_vector(size=LARGE_X)
        med = LARGE_X // 2
        c = nd.slice_axis(a, axis=0, begin=0, end=med)
        assert c.shape[0] == a.shape[0] // 2
        assert c[-1][0] == (med - 1)

    def check_gather():
        arr = mx.nd.ones(LARGE_X)
        # Passing dtype=np.int64 since randomly generated indices are
        # very large that exceeds int32 limits.
        idx = mx.nd.random.randint(0, LARGE_X, 10, dtype=np.int64)
        # Calls gather_nd internally
        tmp = arr[idx]
        assert np.sum(tmp.asnumpy() == 1) == 10
        # Calls gather_nd internally
        arr[idx] += 1
        assert np.sum(arr[idx].asnumpy() == 2) == 10

    def check_infer_shape():
        data_1 = mx.symbol.Variable('data_1')
        data_2 = mx.symbol.Variable('data_2')
        add = data_1+data_2
        # > add.infer_shape(data_1=(LARGE_X,), data_2=(LARGE_X,))
        # OUTPUT - arg_shapes, out_shapes, aux_shapes
        _, out_shapes, _ = add.infer_shape(data_1=(LARGE_X,), data_2=(LARGE_X,))
        assert out_shapes == [(LARGE_X,)]

    def check_astype():
        x = create_vector(size=LARGE_X//4)
        x = nd.tile(x, 4)
        y = x.astype('int32')
        assert y.dtype == np.int32
        assert y[-1] == LARGE_X//4-1

    def check_cast():
        x = create_vector(size=LARGE_X//4)
        x = nd.tile(x, 4)
        y = nd.cast(x, np.int32)
        assert y.dtype == np.int32
        assert y[-1] == LARGE_X//4-1

    def check_load_save():
        x = create_vector(size=LARGE_X)
        with tempfile.TemporaryDirectory() as tmp:
            tmpfile = os.path.join(tmp, 'large_vector')
            nd.save(tmpfile, [x])
            y = nd.load(tmpfile)
            y = y[0]
            assert x[0] == y[0]
            assert x[-1] == y[-1]

    def check_binary_broadcast():
        def check_correctness(mxnet_op, numpy_op, atol=1e-3):
            a = mx.nd.ones(LARGE_X).as_np_ndarray()
            b = 2*mx.nd.ones(LARGE_X).as_np_ndarray()
            res = mxnet_op(a, b)
            np_res = numpy_op(1, 2)
            assert np.abs(res[-1] - np_res) < atol
        check_correctness(mx.np.arctan2, np.arctan2)
        check_correctness(mx.np.hypot, np.hypot)

    check_ndarray_zeros()
    check_ndarray_ones()
    check_ndarray_empty()
    check_ndarray_random_uniform()
    check_ndarray_random_randint()
    check_ndarray_random_exponential()
    check_ndarray_random_gamma()
    check_ndarray_random_generalized_negative_binomial()
    check_ndarray_random_multinomial()
    check_ndarray_random_negative_binomial()
    check_ndarray_random_normal()
    check_ndarray_random_poisson()
    check_ndarray_random_randn()
    check_ndarray_random_shuffle()
    check_full()
    check_repeat()
    check_clip()
    check_slice()
    check_slice_assign()
    check_take()
    check_expand_dims()
    check_squeeze()
    check_broadcast_div()
    check_size()
    check_copy()
    check_copy_to()
    check_zeros_like()
    check_ones_like()
    check_shape()
    check_concat()
    check_slice_like()
    check_slice_axis()
    check_gather()
    check_infer_shape()
    check_astype()
    check_cast()
    check_load_save()
    check_binary_broadcast()


@pytest.mark.timeout(0)
def test_basic():
    def check_elementwise():
        a = nd.ones(shape=LARGE_X)
        b = nd.ones(shape=LARGE_X)
        res = a + b
        assert res[-1].asnumpy() == 2
        res = a + 1
        assert res[-1].asnumpy() == 2
        res = nd.sqrt(a + 8)
        assert res[-1].asnumpy() == 3

    def check_argmin():
        a = create_vector(LARGE_X, dtype=np.float32)
        assert a[0] == 0
        idx = mx.nd.argmin(a, axis=0)
        assert idx[0] == 0
        assert idx.shape[0] == 1

    @pytest.mark.skip(reason="Memory doesn't free up after stacked execution with other ops, " +
                      "tracked at https://github.com/apache/incubator-mxnet/issues/17411")
    def check_argsort():
        a = create_vector(size=LARGE_X)
        s = nd.argsort(a, axis=0, is_ascend=False, dtype=np.int64)
        assert s[0] == (LARGE_X - 1)

    @pytest.mark.skip(reason="Memory doesn't free up after stacked execution with other ops, " +
                      "tracked at https://github.com/apache/incubator-mxnet/issues/17411")
    def check_sort():
        a = create_vector(size=LARGE_X)

        def check_descend(x):
            s = nd.sort(x, axis=0, is_ascend=False)
            assert s[-1] == 0

        def check_ascend(x):
            s = nd.sort(x, is_ascend=True)
            assert s[0] == 0

        check_descend(a)
        check_ascend(a)

    @pytest.mark.skip(reason="Memory doesn't free up after stacked execution with other ops, " +
                      "tracked at https://github.com/apache/incubator-mxnet/issues/17411")
    def check_topk():
        a = create_vector(size=LARGE_X)
        ind = nd.topk(a, k=10, axis=0, dtype=np.int64)
        for i in range(10):
            assert ind[i] == (LARGE_X - i - 1)
        ind, val = mx.nd.topk(a, k=3, axis=0, dtype=np.int64, ret_typ="both", is_ascend=False)
        assert np.all(ind == val)
        val = nd.topk(a, k=1, axis=0, dtype=np.int64, ret_typ="value")
        assert val == (LARGE_X - 1)

    def check_mean():
        a = nd.arange(-LARGE_X // 2, LARGE_X // 2 + 1, dtype=np.int64)
        b = nd.mean(a, axis=0)
        assert b == 0

    def check_exponent_logarithm_operators():
        a = 2*nd.ones(shape=LARGE_X)
        # exponent
        result = nd.exp(a)
        assert result[-1] == 7.389056
        assert result.shape == a.shape
        # exponent minus 1
        result = nd.expm1(a)
        assert result[-1] == 6.389056
        assert result.shape == a.shape
        # log2
        result = nd.log2(a)
        assert result[-1] == 1
        assert result.shape == a.shape
        # log10
        result = nd.log10(a)
        assert result[-1] == 0.30103
        assert result.shape == a.shape
        # log1p
        result = nd.log1p(a)
        assert result[-1] == 1.0986123
        assert result.shape == a.shape
        # log
        result = nd.log(a)
        assert result[-1] == 0.6931472
        assert result.shape == a.shape

    def check_power_operators():
        a = 2*nd.ones(shape=LARGE_X)
        # sqrt
        result = nd.sqrt(a)
        assert result[-1] == 1.4142135
        assert result.shape == a.shape
        # rsqrt
        result = nd.rsqrt(a)
        assert result[-1] == 0.70710677
        assert result.shape == a.shape
        # cbrt
        result = nd.cbrt(a)
        assert result[-1] == 1.2599211
        assert result.shape == a.shape
        # rcbrt
        result = nd.rcbrt(a)
        assert result[-1] == 0.7937005
        assert result.shape == a.shape
        # square
        result = nd.square(a)
        assert result[-1] == 4
        assert result.shape == a.shape
        # reciprocal
        result = nd.reciprocal(a)
        assert result[-1] == 0.5
        assert result.shape == a.shape

    def check_add():
        a = nd.ones(shape=LARGE_X)
        b = nd.ones(shape=LARGE_X)
        c = b
        c = c.__add__(a)
        assert c[-1] == 2
        assert c.shape == a.shape

    def check_sub():
        a = 3*nd.ones(shape=LARGE_X)
        b = nd.ones(shape=LARGE_X)
        c = b
        c = c.__sub__(a)
        assert c[-1] == -2
        assert c.shape == a.shape

    def check_rsub():
        a = 3*nd.ones(shape=LARGE_X)
        b = nd.ones(shape=LARGE_X)
        c = b
        c = c.__rsub__(a)
        assert c[-1] == 2
        assert c.shape == a.shape

    def check_neg():
        a = nd.ones(shape=LARGE_X)
        c = a
        c = c.__neg__()
        assert c[-1] == -1
        assert c.shape == a.shape

    def check_mul():
        a = 2*nd.ones(shape=LARGE_X)
        b = 3*nd.ones(shape=LARGE_X)
        c = b
        c = c.__mul__(a)
        assert c[-1] == 6
        assert c.shape == a.shape

    def check_div():
        a = 2*nd.ones(shape=LARGE_X)
        b = 3*nd.ones(shape=LARGE_X)
        c = b
        c = c.__div__(a)
        assert c[-1] == 3/2
        assert c.shape == a.shape

    def check_rdiv():
        a = 2*nd.ones(shape=LARGE_X)
        b = 3*nd.ones(shape=LARGE_X)
        c = b
        c = c.__rdiv__(a)
        assert c[-1] == 2/3
        assert c.shape == a.shape

    def check_mod():
        a = 2*nd.ones(shape=LARGE_X)
        b = 3*nd.ones(shape=LARGE_X)
        c = b
        c = c.__mod__(a)
        assert c[-1] == 1
        assert c.shape == a.shape

    def check_rmod():
        a = 2*nd.ones(shape=LARGE_X)
        b = 3*nd.ones(shape=LARGE_X)
        c = b
        c = c.__rmod__(a)
        assert c[-1] == 2
        assert c.shape == a.shape

    def check_imod():
        a = 2*nd.ones(shape=LARGE_X)
        b = 3*nd.ones(shape=LARGE_X)
        c = b
        c = c.__imod__(a)
        assert c[-1] == 1
        assert c.shape == a.shape

    def check_pow():
        a = 2*nd.ones(shape=LARGE_X)
        b = 3*nd.ones(shape=LARGE_X)
        c = b
        c = c.__pow__(a)
        assert c[-1] == 9
        assert c.shape == a.shape

    def check_rpow():
        a = 2*nd.ones(shape=LARGE_X)
        b = 3*nd.ones(shape=LARGE_X)
        c = b
        c = c.__rpow__(a)
        assert c[-1] == 8
        assert c.shape == a.shape

    def check_sum():
        a = nd.ones(LARGE_X)
        b = nd.sum(a, axis=0)
        assert b[0] == LARGE_X

    def check_prod():
        a = nd.ones(LARGE_X)
        b = nd.prod(a, axis=0)
        assert b[0] == 1

    def check_min():
        a = create_vector(size=LARGE_X)
        b = nd.min(a, axis=0)
        assert b[0] == 0
        assert b[-1] == 0

    def check_max():
        a = create_vector(size=LARGE_X)
        b = nd.max(a, axis=0)
        assert b[0] == (LARGE_X - 1)

    def check_argmax():
        a = nd.ones(LARGE_X)
        b = nd.zeros(LARGE_X)
        c = nd.concat(a, b, dim=0)
        d = nd.argmax(c, axis=0)
        assert c.shape[0] == (2 * LARGE_X)
        assert d == 0

    def check_iadd():
        a = nd.ones(LARGE_X)
        b = nd.ones(LARGE_X)
        c = b
        c += a
        assert c.shape == a.shape
        assert c[-1] == 2

    def check_isub():
        a = nd.full(LARGE_X, 3)
        b = nd.ones(LARGE_X)
        c = a
        c -= b
        assert c.shape == a.shape
        assert c[-1] == 2

    def check_imul():
        a = nd.full(LARGE_X, 3)
        b = nd.ones(LARGE_X)
        c = b
        c *= a
        assert c.shape == a.shape
        assert c[-1] == 3

    def check_idiv():
        a = nd.full(LARGE_X, 4)
        b = nd.full(LARGE_X, 2)
        c = a
        c /= b
        assert c.shape == a.shape
        assert c[-1] == 2

    def check_eq():
        a = nd.full(LARGE_X, 3)
        b = nd.full(LARGE_X, 3)
        c = (a == b)
        assert (c.asnumpy() == 1).all()

    def check_neq():
        a = nd.full(LARGE_X, 2)
        b = nd.full(LARGE_X, 3)
        c = (a != b)
        assert (c.asnumpy() == 1).all()

    def check_lt():
        a = nd.full(LARGE_X, 2)
        b = nd.full(LARGE_X, 3)
        d = (a <= b)
        assert (d.asnumpy() == 1).all()

    def check_lte():
        a = nd.full(LARGE_X, 2)
        b = nd.full(LARGE_X, 3)
        c = nd.full(LARGE_X, 2)
        d = (a <= b)
        assert (d.asnumpy() == 1).all()
        d = (a <= c)
        assert (d.asnumpy() == 1).all()

    def check_gt():
        a = nd.full(LARGE_X, 3)
        b = nd.full(LARGE_X, 2)
        d = (a > b)
        assert (d.asnumpy() == 1).all()

    def check_gte():
        a = nd.full(LARGE_X, 3)
        b = nd.full(LARGE_X, 2)
        c = nd.full(LARGE_X, 3)
        d = (a >= b)
        assert (d.asnumpy() == 1).all()
        d = (a >= c)
        assert (d.asnumpy() == 1).all()

    def check_logical():
        def check_logical_and(a, b):
            mx_res = mx.nd.logical_and(a, b)
            assert_almost_equal(mx_res[-1].asnumpy(), np.logical_and(a[-1].asnumpy(), b[-1].asnumpy()))

        def check_logical_or(a, b):
            mx_res = mx.nd.logical_or(a, b)
            assert_almost_equal(mx_res[-1].asnumpy(), np.logical_or(a[-1].asnumpy(), b[-1].asnumpy()))

        def check_logical_not(a, b):
            mx_res = mx.nd.logical_not(a, b)
            assert_almost_equal(mx_res[-1].asnumpy(), np.logical_not(a[-1].asnumpy(), b[-1].asnumpy()))

        def check_logical_xor(a, b):
            mx_res = mx.nd.logical_xor(a, b)
            assert_almost_equal(mx_res[-1].asnumpy(), np.logical_xor(a[-1].asnumpy(), b[-1].asnumpy()))

        a = mx.nd.ones(LARGE_X)
        b = mx.nd.zeros(LARGE_X)
        check_logical_and(a, b)
        check_logical_or(a, b)
        check_logical_not(a, b)
        check_logical_xor(a, b)

    def create_input_for_rounding_ops():
        # Creates an vector with values (-LARGE/2 .... -2, -1, 0, 1, 2, .... , LARGE/2-1)
        # then divides each element by 2 i.e (-LARGE/4 .... -1, -0.5, 0, 0.5, 1, .... , LARGE/4-1)
        inp = nd.arange(-LARGE_X//2, LARGE_X//2, dtype=np.float64)
        inp = inp/2
        return inp

    def assert_correctness_of_rounding_ops(output, mid, expected_vals):
        # checks verifies 5 values at the middle positions of the input vector
        # i.e mid-2, mid-1, mid, mid+1, mid+2
        output_idx_to_inspect = [mid-2, mid-1, mid, mid+1, mid+2]
        for i in range(len(output_idx_to_inspect)):
            assert output[output_idx_to_inspect[i]] == expected_vals[i]

    def check_rounding_ops():
        x = create_input_for_rounding_ops()

        def check_ceil():
            y = nd.ceil(x)
            # expected ouput for middle 5 values after applying ceil()
            expected_output = [-1, 0, 0, 1, 1]
            assert_correctness_of_rounding_ops(y, LARGE_X//2, expected_output)

        def check_fix():
            y = nd.fix(x)
            # expected ouput for middle 5 values after applying fix()
            expected_output = [-1, 0, 0, 0, 1]
            assert_correctness_of_rounding_ops(y, LARGE_X//2, expected_output)

        def check_floor():
            y = nd.floor(x)
            # expected ouput for middle 5 values after applying floor()
            expected_output = [-1, -1, 0, 0, 1]
            assert_correctness_of_rounding_ops(y, LARGE_X//2, expected_output)

        def check_rint():
            y = nd.rint(x)
            # expected ouput for middle 5 values after applying rint()
            expected_output = [-1, -1, 0, 0, 1]
            assert_correctness_of_rounding_ops(y, LARGE_X//2, expected_output)

        def check_round():
            y = nd.round(x)
            # expected ouput for middle 5 values after applying round()
            expected_output = [-1, -1, 0, 1, 1]
            assert_correctness_of_rounding_ops(y, LARGE_X//2, expected_output)

        def check_trunc():
            y = nd.trunc(x)
            # expected ouput for middle 5 values after applying trunc()
            expected_output = [-1, 0, 0, 0, 1]
            assert_correctness_of_rounding_ops(y, LARGE_X//2, expected_output)

        check_ceil()
        check_fix()
        check_floor()
        check_rint()
        check_round()
        check_trunc()

    def create_input_for_trigonometric_ops(vals):
        # Creates large vector input of size(LARGE_X) from vals using tile operator
        inp = nd.array(vals)
        inp = nd.tile(inp, LARGE_X//len(vals))
        return inp

    def assert_correctness_of_trigonometric_ops(output, expected_vals):
        # checks verifies 5 values at positions(0, 1, -3, -2, -1) of the input vector
        output_idx_to_inspect = [0, 1, -3, -2, -1]
        for i in range(len(output_idx_to_inspect)):
            assert np.abs(output[output_idx_to_inspect[i]].asnumpy()-expected_vals[i]) <= 1e-3

    def check_trigonometric_ops():
        def check_arcsin():
            x = create_input_for_trigonometric_ops([-1, -.707, 0, .707, 1])
            y = nd.arcsin(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying arcsin()
            expected_output = [-np.pi/2, -np.pi/4, 0, np.pi/4, np.pi/2]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_arccos():
            x = create_input_for_trigonometric_ops([-1, -.707, 0, .707, 1])
            y = nd.arccos(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying arccos()
            expected_output = [np.pi, 3*np.pi/4, np.pi/2, np.pi/4, 0]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_arctan():
            x = create_input_for_trigonometric_ops([-np.Inf, -1, 0, 1, np.Inf])
            y = nd.arctan(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying arctan()
            expected_output = [-np.pi/2, -np.pi/4, 0, np.pi/4, np.pi/2]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_sin():
            x = create_input_for_trigonometric_ops([-np.pi/2, -np.pi/4, 0, np.pi/4, np.pi/2])
            y = nd.sin(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying sin()
            expected_output = [-1, -.707, 0, .707, 1]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_cos():
            x = create_input_for_trigonometric_ops([0, np.pi/4, np.pi/2, 3*np.pi/4, np.pi])
            y = nd.cos(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying cos()
            expected_output = [1, .707, 0, -.707, -1]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_tan():
            x = create_input_for_trigonometric_ops([-np.pi/6, -np.pi/4, 0, np.pi/4, np.pi/6])
            y = nd.tan(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying tan()
            expected_output = [-.577, -1, 0, 1, .577]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_arcsinh():
            x = create_input_for_trigonometric_ops([-np.pi/2, -np.pi/4, 0, np.pi/4, np.pi/2])
            y = nd.arcsinh(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying arcsinh()
            expected_output = [np.arcsinh(-np.pi/2), np.arcsinh(-np.pi/4), 0, np.arcsinh(np.pi/4), np.arcsinh(np.pi/2)]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_arccosh():
            x = create_input_for_trigonometric_ops([1, np.pi/2, 3*np.pi/4, np.pi, 5*np.pi/4])
            y = nd.arccosh(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying arccosh()
            expected_output = [0, np.arccosh(np.pi/2), np.arccosh(3*np.pi/4), np.arccosh(np.pi), np.arccosh(5*np.pi/4)]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_arctanh():
            x = create_input_for_trigonometric_ops([-1/4, -1/2, 0, 1/4, 1/2])
            y = nd.arctanh(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying arctanh()
            expected_output = [np.arctanh(-1/4), np.arctanh(-1/2), 0, np.arctanh(1/4), np.arctanh(1/2)]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_sinh():
            x = create_input_for_trigonometric_ops([-np.pi/2, -np.pi/4, 0, np.pi/4, np.pi/2])
            y = nd.sinh(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying sinh()
            expected_output = [np.sinh(-np.pi/2), np.sinh(-np.pi/4), 0, np.sinh(np.pi/4), np.sinh(np.pi/2)]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_cosh():
            x = create_input_for_trigonometric_ops([0, 1, np.pi/2, 3*np.pi/4, np.pi])
            y = nd.cosh(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying cosh()
            expected_output = [1, np.cosh(1), np.cosh(np.pi/2), np.cosh(3*np.pi/4), np.cosh(np.pi)]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_tanh():
            x = create_input_for_trigonometric_ops([-1/4, -1/2, 0, 1/4, 1/2])
            y = nd.tanh(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying tanh()
            expected_output = [np.tanh(-1/4), np.tanh(-1/2), 0, np.tanh(1/4), np.tanh(1/2)]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_radians():
            x = create_input_for_trigonometric_ops([0, 90, 180, 270, 360])
            y = nd.radians(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying radians()
            expected_output = [0, np.pi/2, np.pi, 3*np.pi/2, 2*np.pi]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        def check_degrees():
            x = create_input_for_trigonometric_ops([0, np.pi/2, np.pi, 3*np.pi/2, 2*np.pi])
            y = nd.degrees(x)
            # expected ouput for indices=(0, 1, -3, -2, -1) after applying degrees()
            expected_output = [0, 90, 180, 270, 360]
            assert_correctness_of_trigonometric_ops(y, expected_output)

        check_arcsin()
        check_arccos()
        check_arctan()
        check_sin()
        check_cos()
        check_tan()
        check_arcsinh()
        check_arccosh()
        check_arctanh()
        check_sinh()
        check_cosh()
        check_tanh()
        check_radians()
        check_degrees()

    def check_add_n():
        x = [nd.ones(LARGE_X)]
        y = nd.add_n(*x)
        assert y[0] == 1
        assert y[-1] == 1

    def check_modulo():
        x = mx.nd.ones(LARGE_X)*6
        y = mx.nd.ones(LARGE_X)*4
        z = (x % y)
        assert z[0] == 2
        assert z[-1] == 2
        x = mx.nd.ones(LARGE_X)*5
        z = nd.modulo(x, y)
        assert z[0] == 1
        assert z[-1] == 1

    def check_maximum():
        x = mx.nd.ones(LARGE_X)*3
        y = mx.nd.ones(LARGE_X)*4
        z = nd.maximum(x, y)
        assert z[0] == 4
        assert z[-1] == 4
        z = nd.maximum(x, 5)
        assert z[0] == 5
        assert z[-1] == 5

    def check_minimum():
        x = mx.nd.ones(LARGE_X)*3
        y = mx.nd.ones(LARGE_X)*2
        z = nd.minimum(x, y)
        assert z[0] == 2
        assert z[-1] == 2
        z = nd.minimum(x, 5)
        assert z[0] == 3
        assert z[-1] == 3

    check_elementwise()
    check_argmin()
    check_argsort()
    check_sort()
    check_topk()
    check_mean()
    check_exponent_logarithm_operators()
    check_power_operators()
    check_add()
    check_sub()
    check_rsub()
    check_neg()
    check_mul()
    check_div()
    check_rdiv()
    check_mod()
    check_rmod()
    check_imod()
    check_pow()
    check_rpow()
    check_sum()
    check_prod()
    check_min()
    check_max()
    check_argmax()
    check_iadd()
    check_isub()
    check_imul()
    check_idiv()
    check_eq()
    check_neq()
    check_lt()
    check_lte()
    check_gt()
    check_gte()
    check_logical()
    check_rounding_ops()
    check_trigonometric_ops()
    check_add_n()
    check_modulo()
    check_maximum()
    check_minimum()