code: def test_nn():
      def check_gluon_embedding():
          m = gluon.nn.Embedding(SMALL_Y, MEDIUM_X)
          m.initialize()
          a = nd.zeros((MEDIUM_X, SMALL_Y))
          b = m(a)
          assert b.shape == (MEDIUM_X, SMALL_Y, MEDIUM_X)
          assert b.asnumpy().size == LARGE_SIZE
  
      def check_fully_connected():
          a = nd.ones(shape=(LARGE_X, SMALL_Y))
          b = nd.ones(shape=(SMALL_Y, SMALL_Y))
          c = nd.ones(shape=(b.shape[0],))
  
          # w/o bias
          res = nd.FullyConnected(
              a, b, num_hidden=b.shape[0], no_bias=True
          )
          assert np.sum(res[-1].asnumpy() == a.shape[1]) == b.shape[0]
  
          # w/ bias
          res = nd.FullyConnected(
              a, b, c, num_hidden=b.shape[0], no_bias=False
          )
          assert (
              np.sum(res[-1].asnumpy() == a.shape[1] + 1) == b.shape[0]
          )
  
      def check_dense():
          data = mx.nd.ones(shape=(50 * 1000 * 1000, 100))
          linear = gluon.nn.Dense(100)
          linear.initialize()
          res = linear(data)
          assert res.shape == (50000000, 100)
  
      def check_softmax():
          input_data = mx.nd.ones((SMALL_Y, LARGE_X))
          for axis in [0, 1]:
              true_output = np.full(
                  (SMALL_Y, LARGE_X), (1 / input_data.shape[axis])
              )
              output = nd.softmax(input_data, axis=axis)
              assert_almost_equal(
                  output.asnumpy(), true_output, rtol=1e-5, atol=1e-5
              )
  
      def check_softmax_cross_entropy():
          # dtype of input data, mxnet cross entropy set explicitly to float64
          # numpy implicitly takes care of double precision
          batch_size = SMALL_Y
          num_labels = LARGE_X
          input_data = mx.nd.ones(
              (batch_size, num_labels), dtype="float64"
          )
          input_label = mx.nd.zeros((batch_size,), dtype="float64")
          true_softmax = np.full(
              (batch_size, num_labels), (1 / num_labels)
          )
          # use 1/batch_size when softmax axis=0
          # here 1/num_labels since softmax_cross_entropy uses default axis
          # by default axis=1
          np_one_hot_label = np.zeros((batch_size, num_labels))
          np_one_hot_label[:, 0] = 1
          true_softmax_cross_entropy = np.sum(
              -np.log(true_softmax) * np_one_hot_label
          )
          mx_softmax_cross_entropy = mx.nd.softmax_cross_entropy(
              input_data, input_label, dtype="float64"
          )
          assert_almost_equal(
              mx_softmax_cross_entropy.asnumpy(),
              true_softmax_cross_entropy,
              rtol=1e-3,
              atol=1e-5,
          )
  
      def check_softmax_activation():
          data = nd.random_normal(shape=(2**29, 2, 2, 2))
          out = nd.random_normal(shape=(2**29, 2, 2, 2))
  
          res = nd.SoftmaxActivation(data=data, out=out)
  
          assert res.shape[0] == 536870912
          assert res.shape[1] == 2
          assert res.shape[2] == 2
          assert res.shape[3] == 2
  
      def np_softmax(x, axis=-1, temperature=1.0):
          x = x - np.max(x, axis=axis, keepdims=True)
          x = np.exp(x / temperature)
          x /= np.sum(x, axis=axis, keepdims=True)
          return x
  
      @pytest.mark.skip(
          reason="log_softmax flaky, tracked at "
          "https://github.com/apache/mxnet/issues/17397"
      )
      def check_log_softmax():
          ndim = 2
          shape = (SMALL_Y, LARGE_X)
          axis = np.random.randint(0, ndim)
          data = np.random.uniform(-2, 2, size=shape)
          sym = mx.sym.log_softmax(axis=axis - ndim)
          check_symbolic_forward(
              sym, [data], [np.log(np_softmax(data, axis=axis) + 1e-20)]
          )
  
      # TODO: correctness of prelu (currently flaky)
      def check_leaky_relu():
          a = -1 * mx.nd.ones((LARGE_X, SMALL_Y))
  
          def check_leaky():
              res = mx.nd.LeakyReLU(a, act_type="leaky", slope=0.3)
              assert_almost_equal(
                  res[-1][-1].asnumpy(),
                  0.3 * a[-1][-1].asnumpy(),
                  atol=1e-3,
                  rtol=1e-3,
              )
  
          def check_elu():
              res = mx.nd.LeakyReLU(a, act_type="elu", slope=0.3)
              assert_almost_equal(
                  res[-1][-1].asnumpy(),
                  0.3 * (np.exp(a[-1][-1].asnumpy()) - 1),
                  atol=1e-3,
                  rtol=1e-3,
              )
  
          def check_selu():
              lam = 1.0507009873554804934193349852946
              alpha = 1.6732632423543772848170429916717
              res = mx.nd.LeakyReLU(a, act_type="selu")
              assert_almost_equal(
                  res[-1][-1].asnumpy(),
                  (lam * alpha * (np.exp(a[-1][-1].asnumpy()) - 1)),
                  atol=1e-3,
                  rtol=1e-3,
              )
  
          def check_rrelu():
              lower = 0.125
              upper = 0.333999991
              res = mx.nd.LeakyReLU(a, act_type="rrelu")
              assert_almost_equal(
                  res[0][-1][-1].asnumpy(),
                  (lower + upper) / 2 * a[-1][-1].asnumpy(),
                  atol=1e-3,
                  rtol=1e-3,
              )
  
          check_leaky()
          check_elu()
          check_selu()
          check_rrelu()
  
      def check_pooling():
          a = mx.nd.ones((MEDIUM_X, 200, SMALL_Y, SMALL_Y))
  
          def check_avg_pooling():
              res = mx.nd.Pooling(a, kernel=(5, 5), pool_type="avg")
              assert_almost_equal(
                  res[-1][-1][-1][-1].asnumpy(),
                  1.0000001,
                  atol=1e-3,
                  rtol=1e-3,
              )
              assert res.shape[-1] == SMALL_Y - 5 + 1
  
          def check_max_pooling():
              res = mx.nd.Pooling(a, kernel=(5, 5), pool_type="max")
              assert_almost_equal(
                  res[-1][-1][-1][-1].asnumpy(),
                  1.0,
                  atol=1e-3,
                  rtol=1e-3,
              )
              assert res.shape[-1] == SMALL_Y - 5 + 1
  
          def check_sum_pooling():
              res = mx.nd.Pooling(a, kernel=(5, 5), pool_type="sum")
              assert_almost_equal(
                  res[-1][-1][-1][-1].asnumpy(),
                  25,
                  atol=1e-3,
                  rtol=1e-3,
              )
              assert res.shape[-1] == SMALL_Y - 5 + 1
  
          def check_lp_pooling():
              res = mx.nd.Pooling(
                  a, kernel=(5, 5), pool_type="lp", p_value=2
              )
              assert_almost_equal(
                  res[-1][-1][-1][-1].asnumpy(),
                  5.0,
                  atol=1e-3,
                  rtol=1e-3,
              )
              assert res.shape[-1] == SMALL_Y - 5 + 1
  
              res = mx.nd.Pooling(
                  a, kernel=(5, 5), pool_type="lp", p_value=1
              )
              assert_almost_equal(
                  res[-1][-1][-1][-1].asnumpy(),
                  25.0,
                  atol=1e-3,
                  rtol=1e-3,
              )
              assert res.shape[-1] == SMALL_Y - 5 + 1
  
          check_avg_pooling()
          check_max_pooling()
          check_sum_pooling()
          check_lp_pooling()
  
      def check_layer_norm():
          dtype = np.float32
          forward_check_eps = 1e-3
          axis = 1
          eps = 1e-5
          in_shape = (LARGE_X, SMALL_Y)
          ctx = mx.cpu()
  
          def npy_layer_norm(data, gamma, beta, axis=1, eps=1e-5):
              if axis < 0:
                  axis += data.ndim
              broadcast_shape = [1 for _ in range(data.ndim)]
              broadcast_shape[axis] = data.shape[axis]
              mean = data.mean(axis=axis, keepdims=True).astype(dtype)
              var = data.var(axis=axis, keepdims=True).astype(dtype)
              std = np.sqrt(var + dtype(eps)).astype(dtype)
              out = np.reshape(gamma, broadcast_shape) * (
                  data - mean
              ) / std + np.reshape(beta, broadcast_shape)
              return out
  
          data = np.random.normal(0, 1, in_shape).astype(dtype)
          gamma = np.random.normal(0, 1, (in_shape[axis],)).astype(
              dtype
          )
          beta = np.random.normal(0, 1, (in_shape[axis],)).astype(dtype)
          data_s = mx.symbol.Variable("data")
          gamma_s = mx.symbol.Variable("gamma")
          beta_s = mx.symbol.Variable("beta")
          out_s = mx.symbol.LayerNorm(
              data=data_s,
              gamma=gamma_s,
              beta=beta_s,
              axis=axis,
              eps=eps,
          )
          exe = out_s._simple_bind(ctx, data=in_shape)
          exe.arg_dict["data"][:] = data
          exe.arg_dict["gamma"][:] = gamma
          exe.arg_dict["beta"][:] = beta
          out_nd = exe.forward()[0]
          out = npy_layer_norm(data, gamma, beta, axis, eps)
          assert_almost_equal(
              out,
              out_nd.asnumpy(),
              forward_check_eps,
              forward_check_eps,
          )
  
      # TODO: correctness of dropout
      # currently only test for dropout to work
      # since testing for correctness involves flakiness issue #14288
      def check_dropout():
          shape = (LARGE_X, SMALL_Y)
          x = mx.sym.var("data")
          y = mx.sym.Dropout(x, p=1, cudnn_off=True)
          exe = y._simple_bind(ctx=default_device(), data=shape)
          exe.arg_arrays[0][:] = 1
          out = exe.forward(is_train=True)
          nd.waitall()
          assert out[0].shape == shape
  
      def check_activation():
          x = mx.nd.ones((LARGE_X, SMALL_Y))
          check_x = -2
          x[-1, -1] = check_x
          # Hyperbolic tangent (tanh)
          # y = (exp(x)-exp(-x))/(exp(x)+exp(-x))
          y = mx.nd.Activation(x, act_type="tanh")
          tanh_x = (np.exp(check_x) - np.exp(-check_x)) / (
              np.exp(check_x) + np.exp(-check_x)
          )
          assert y[-1][-1] == np.float32(tanh_x)
          # Recitified Linear Unit (relu)
          # y = max(x,0)
          y = mx.nd.Activation(x, act_type="relu")
          assert y[-1][-1] == 0
          # Sigmoid
          # y = x/(1+abs(x))
          y = mx.nd.Activation(x, act_type="sigmoid")
          sigmoid_x = 1 / (1 + math.exp(-check_x))
          assert_almost_equal(
              y[-1][-1].asnumpy(),
              np.float32(sigmoid_x),
              atol=1e-3,
              rtol=1e-3,
          )
          # Soft Sign
          # y = 1/(1+exp(-x))
          y = mx.nd.Activation(x, act_type="softsign")
          softsign_x = check_x / (1 + abs(check_x))
          assert y[-1][-1] == np.float32(softsign_x)
  
      # TODO: correctness of batchnorm
      # in future, we could test if mean, var of output
      # matches target output's mean, var
      def check_batchnorm():
          def get_np_mean_var(
              data,
              running_mean,
              running_var,
              eps,
              use_global_status=True,
          ):
              if not use_global_status:
                  # train mode, calculate the real mean and var
                  mean = np.mean(data, axis=(0, 2, 3))
                  mean_broad = np.expand_dims(mean, axis=0)
                  mean_broad = np.expand_dims(mean_broad, axis=2)
                  mean_broad = np.expand_dims(mean_broad, axis=3)
                  mean_broad = np.broadcast_to(mean_broad, data.shape)
                  var = np.square(data - mean_broad)
                  var = np.mean(var, axis=(0, 2, 3))
              else:
                  # inference mode, use running_mean and running_var instead
                  mean = np.full((data.shape[1],), running_mean)
                  var = np.full((data.shape[1],), running_var)
              # calculate the inverse of standard variance
              invstdvar = 1.0 / np.sqrt(var + eps)
              return mean, invstdvar
  
          # Here use 4D input to cover dnnl BN and non-dnnl BN
          shape = (1, 2, LARGE_X, SMALL_Y)
          axis = 1  # default
          eps = 1e-3
          nch = shape[axis]
          data = mx.nd.ones(shape=shape)
          bn_gamma = mx.nd.random.uniform(shape=(nch,))
          bn_beta = mx.nd.random.uniform(shape=(nch,))
          bn_running_mean = mx.nd.zeros(nch)
          bn_running_var = mx.nd.ones(nch)
          output = mx.nd.BatchNorm(
              data,
              bn_gamma,
              bn_beta,
              bn_running_mean,
              bn_running_var,
              output_mean_var=True,
          )
          assert output[0].shape == shape
          mean, invstdvar = output[1], output[2]
          np_mean, np_invstdvar = get_np_mean_var(
              data.asnumpy(),
              bn_running_mean.asnumpy(),
              bn_running_var.asnumpy(),
              eps,
              use_global_status=True,
          )
          assert_almost_equal(mean.asnumpy(), np_mean)
          assert_almost_equal(invstdvar.asnumpy(), np_invstdvar)
  
      def check_relu():
          def frelu(x):
              return np.maximum(x, 0.0)
  
          def frelu_grad(x):
              return 1.0 * (x > 0.0)
  
          shape = (SMALL_Y, LARGE_X)
          x = mx.symbol.Variable("x")
          y = mx.sym.relu(x)
          xa = np.random.uniform(low=-1.0, high=1.0, size=shape)
          eps = 1e-4
          xa[abs(xa) < eps] = 1.0
          ya = frelu(xa)
          ga = frelu_grad(xa)
          check_symbolic_forward(y, [xa], [ya])
  
      def check_sigmoid():
          def fsigmoid(a):
              return np.divide(1.0, (1.0 + np.exp(-a)))
  
          shape = (SMALL_Y, LARGE_X)
          x = mx.symbol.Variable("x")
          y = mx.sym.sigmoid(x)
          xa = np.random.uniform(low=-1.0, high=1.0, size=shape)
          ya = fsigmoid(xa)
          check_symbolic_forward(y, [xa], [ya])
  
      def check_l2_normalization():
          x = nd.ones((2, LARGE_X * 2))
          x[0] = 3
          x[1] = 4
          # Channel Mode
          z = x.reshape(1, 2, LARGE_X * 2)
          y = nd.L2Normalization(z, mode="channel")
          assert y[0][0][0] == 0.6
          assert y[0][0][-1] == 0.6
          assert y[0][1][0] == 0.8
          assert y[0][1][-1] == 0.8
          # Instance Mode
          z = x.T
          y = nd.L2Normalization(z, mode="instance")
          assert y[0][0] == 0.6
          assert y[0][1] == 0.8
          assert y[-1][0] == 0.6
          assert y[-1][1] == 0.8
          # Spatial Mode
          z = z.reshape(1, 200000000, 2)
          y = nd.L2Normalization(z, mode="spatial")
          assert y[0][0][0] == 0.6
          assert y[0][0][1] == 0.8
          assert y[0][-1][0] == 0.6
          assert y[0][-1][1] == 0.8
  
      def check_instance_norm():
          dtype = np.float32
          forward_check_eps = 1e-3
          axis = -1
          eps = 1e-5
          in_shape = (LARGE_X, 1, SMALL_Y)
          ctx = mx.cpu()
  
          # Implementation of instance normalization using numpy
          def npy_instance_norm(data, gamma, beta, axis, eps=1e-5):
              if axis < 0:
                  axis += data.ndim
              broadcast_shape = [1 for _ in range(data.ndim)]
              broadcast_shape[axis] = data.shape[axis]
              mean = data.mean(axis=axis, keepdims=True).astype(dtype)
              var = data.var(axis=axis, keepdims=True).astype(dtype)
              std = np.sqrt(var + dtype(eps)).astype(dtype)
              out = gamma * (data - mean) / std + beta
              return out
  
          data = np.random.normal(0, 1, in_shape).astype(dtype)
          gamma = np.random.normal(0, 1, (1,)).astype(dtype)
          beta = np.random.normal(0, 1, (1,)).astype(dtype)
          data_s = mx.symbol.Variable("data")
          gamma_s = mx.symbol.Variable("gamma")
          beta_s = mx.symbol.Variable("beta")
          out_s = mx.symbol.InstanceNorm(
              data=data_s, gamma=gamma_s, beta=beta_s, eps=eps
          )
          exe = out_s._simple_bind(ctx, data=in_shape)
          exe.arg_dict["data"][:] = data
          exe.arg_dict["gamma"][:] = gamma
          exe.arg_dict["beta"][:] = beta
          out_nd = exe.forward()[0]
          # Calls implementation of instance norm in numpy and compares the output
          out = npy_instance_norm(data, gamma, beta, axis, eps)
          assert_almost_equal(
              out,
              out_nd.asnumpy(),
              forward_check_eps,
              forward_check_eps,
          )
  
      def check_col2im():
          data = nd.random_normal(shape=(1, 2**30, 4))
          output_size = (2, 2, 1)
          kernel = (1, 1, 1)
  
          res = nd.col2im(
              data=data, output_size=output_size, kernel=kernel
          )
  
          assert res.shape[0] == 1
          assert res.shape[1] == 1073741824
          assert res.shape[2] == 2
          assert res.shape[3] == 2
          assert res.shape[4] == 1
  
      def check_embedding():
          data = nd.random_normal(shape=(LARGE_TENSOR_SHAPE, 1))
          weight = nd.random_normal(shape=(LARGE_TENSOR_SHAPE, 1))
          input_dim = LARGE_TENSOR_SHAPE
          output_dim = 1
  
          out = nd.Embedding(
              data=data,
              weight=weight,
              input_dim=input_dim,
              output_dim=output_dim,
          )
  
          assert out.shape[0] == LARGE_TENSOR_SHAPE
          assert out.shape[1] == 1
  
      def check_spatial_transformer():
          data = nd.random_normal(shape=(2, 2**29, 1, 6))
          loc = nd.random_normal(shape=(2, 6))
          transform_type = "affine"
          sampler_type = "bilinear"
          target_shape = (2, 6)
  
          res = nd.SpatialTransformer(
              data=data,
              loc=loc,
              transform_type=transform_type,
              sampler_type=sampler_type,
              target_shape=target_shape,
          )
  
          assert res.shape[0] == 2
          assert res.shape[1] == 536870912
          assert res.shape[2] == 2
          assert res.shape[3] == 6
  
      def check_ravel():
          data = nd.random_normal(shape=(2, LARGE_TENSOR_SHAPE))
          shape = (2, 10)
  
          out = nd.ravel_multi_index(data=data, shape=shape)
  
          assert out.shape[0] == LARGE_TENSOR_SHAPE
  
      def check_cumsum():
          a = nd.ones((LARGE_X, SMALL_Y))
          axis = 1
  
          res = nd.cumsum(a=a, axis=axis)
  
          assert res.shape[0] == LARGE_X
          assert res.shape[1] == SMALL_Y
          assert res[0][SMALL_Y - 1] == 50.0
  
      def check_multi_lars():
          lrs = nd.random_normal(shape=(LARGE_TENSOR_SHAPE + 1, 1))
          weights_sum_sq = nd.random_normal(
              shape=(LARGE_TENSOR_SHAPE + 1, 1)
          )
          grads_sum_sq = nd.random_normal(
              shape=(LARGE_TENSOR_SHAPE + 1, 1)
          )
          wds = nd.random_normal(shape=(LARGE_TENSOR_SHAPE + 1, 1))
          eta = 0.1
          eps = 0.9
  
          out = nd.multi_lars(
              lrs=lrs,
              weights_sum_sq=weights_sum_sq,
              grads_sum_sq=grads_sum_sq,
              wds=wds,
              eta=eta,
              eps=eps,
          )
  
          assert out.shape[0] == LARGE_TENSOR_SHAPE + 1
          assert out.shape[1] == 1
  
          # Trigger lazy evaluation of the output NDArray and ensure that it has been filled
          assert type(out[0, 0].asscalar()).__name__ == "float32"
  
      def check_rnn():
          data = nd.random_normal(shape=(RNN_LARGE_TENSOR, 4, 4))
          parameters_relu_tanh = nd.random_normal(shape=(7,))
          parameters_lstm = nd.random_normal(shape=(28,))
          parameters_gru = nd.random_normal(shape=(21,))
          state = nd.random_normal(shape=(1, 4, 1))
          state_cell = nd.random_normal(shape=(1, 4, 1))
          mode_relu = "rnn_relu"
          mode_tanh = "rnn_tanh"
          mode_lstm = "lstm"
          mode_gru = "gru"
          state_size = 1
          num_layers = 1
  
          out_relu = nd.RNN(
              data=data,
              parameters=parameters_relu_tanh,
              state=state,
              mode=mode_relu,
              state_size=state_size,
              num_layers=num_layers,
          )
  
          out_tanh = nd.RNN(
              data=data,
              parameters=parameters_relu_tanh,
              state=state,
              mode=mode_tanh,
              state_size=state_size,
              num_layers=num_layers,
          )
  
          out_lstm = nd.RNN(
              data=data,
              parameters=parameters_lstm,
              state=state,
              mode=mode_lstm,
              state_cell=state_cell,
              state_size=state_size,
              num_layers=num_layers,
          )
  
          out_gru = nd.RNN(
              data=data,
              parameters=parameters_gru,
              state=state,
              mode=mode_gru,
              state_size=state_size,
              num_layers=num_layers,
          )
  
          for out in [out_relu, out_tanh, out_lstm, out_gru]:
              assert out.shape[0] == RNN_LARGE_TENSOR
              assert out.shape[1] == 4
              assert out.shape[2] == 1
  
              assert type(out[0, 0, 0].asscalar()).__name__ == "float32"
  
      check_gluon_embedding()
      check_fully_connected()
      check_dense()
      check_softmax()
      check_softmax_cross_entropy()
      check_softmax_activation()
      check_log_softmax()
      check_leaky_relu()
      check_pooling()
      check_layer_norm()
      check_dropout()
      check_activation()
      check_batchnorm()
      check_relu()
      check_sigmoid()
      check_l2_normalization()
      check_instance_norm()
      check_col2im()
      check_embedding()
      check_spatial_transformer()
      check_ravel()
      check_cumsum()
      check_multi_lars()
      check_rnn()
  
file_path: mxnet/tests/nightly/test_large_array.py
imports: import os
  import sys
  import tempfile
  import math
  import numpy
  import mxnet
  from mxnet.test_utils import rand_ndarray, assert_almost_equal, rand_coord_2d, default_device, check_symbolic_forward, create_2d_tensor
  from mxnet.util import TemporaryDirectory
  from mxnet import gluon, nd
  from common import with_seed
  import pytest
label_desc: Use image node_id 1 for calling a global function with example usage: test_nn() and returns: x, out, mean,  invstdvar, np, unknown, np, out