code: def test_pad_wrap_batchify():
      def _verify_padded_arr(
          padded_arr, original_arr, pad_axis, pad_val, pad_length, dtype
      ):
          ndim = original_arr.ndim
          slices_data = [slice(None) for _ in range(ndim)]
          slices_data[pad_axis] = slice(original_arr.shape[axis])
          assert_allclose(padded_arr[tuple(slices_data)], original_arr)
          if original_arr.shape[pad_axis] < pad_length:
              slices_pad_val = [slice(None) for _ in range(ndim)]
              slices_pad_val[axis] = slice(
                  original_arr.shape[pad_axis], None
              )
              pad_val_in_arr = padded_arr[tuple(slices_pad_val)]
              assert_allclose(
                  pad_val_in_arr,
                  (np.ones_like(pad_val_in_arr) * pad_val).astype(
                      dtype
                  ),
              )
  
      batch_size = 6
      for ndim in range(1, 3):
          for axis in range(-ndim, ndim):
              for length_min, length_max in [(3, 4), (3, 7)]:
                  for pad_val in [-1, 0]:
                      for dtype in [
                          np.uint8,
                          np.int32,
                          np.int64,
                          np.float16,
                          np.float32,
                          np.float64,
                      ]:
                          # Each instance contains a single array
                          for _dtype in [None, dtype]:
                              shapes = [
                                  [2 for _ in range(ndim)]
                                  for _ in range(batch_size)
                              ]
                              for i in range(len(shapes)):
                                  shapes[i][axis] = np.random.randint(
                                      length_min, length_max
                                  )
                              random_data_npy = [
                                  np.random.normal(0, 1, shape).astype(
                                      dtype
                                  )
                                  for shape in shapes
                              ]
                              batchify_fn = batchify.Pad(
                                  axis=axis, val=pad_val, dtype=_dtype
                              )
                              batch_data = batchify_fn(random_data_npy)
                              batch_data_use_mx = batchify_fn(
                                  [
                                      mx.np.array(ele, dtype=dtype)
                                      for ele in random_data_npy
                                  ]
                              )
                              assert_allclose(
                                  batch_data_use_mx.asnumpy(),
                                  batch_data.asnumpy(),
                              )
                              assert (
                                  batch_data.dtype
                                  == batch_data_use_mx.dtype
                                  == dtype
                              )
                              batch_data = batch_data.asnumpy()
                              for i in range(batch_size):
                                  pad_length = max(
                                      shape[axis] for shape in shapes
                                  )
                                  _verify_padded_arr(
                                      batch_data[i],
                                      random_data_npy[i],
                                      axis,
                                      pad_val,
                                      pad_length,
                                      dtype,
                                  )
                              # Each instance contains 3 arrays, we pad part of them according to index
                              TOTAL_ELE_NUM = 3
                              for pad_index in [
                                  [0],
                                  [1],
                                  [2],
                                  [0, 1],
                                  [1, 2],
                                  [0, 1, 2],
                              ]:
                                  shapes = [
                                      [
                                          [2 for _ in range(ndim)]
                                          for _ in range(batch_size)
                                      ]
                                      for _ in range(TOTAL_ELE_NUM)
                                  ]
                                  for j in pad_index:
                                      for i in range(batch_size):
                                          shapes[j][i][
                                              axis
                                          ] = np.random.randint(
                                              length_min, length_max
                                          )
                                  random_data_npy = [
                                      tuple(
                                          np.random.normal(
                                              0, 1, shapes[j][i]
                                          ).astype(dtype)
                                          for j in range(TOTAL_ELE_NUM)
                                      )
                                      for i in range(batch_size)
                                  ]
                                  batchify_fn = []
                                  for j in range(TOTAL_ELE_NUM):
                                      if j in pad_index:
                                          batchify_fn.append(
                                              batchify.Pad(
                                                  axis=axis,
                                                  val=pad_val,
                                                  dtype=_dtype,
                                              )
                                          )
                                      else:
                                          batchify_fn.append(
                                              batchify.Stack(
                                                  dtype=_dtype
                                              )
                                          )
                                  batchify_fn = batchify.Tuple(
                                      batchify_fn
                                  )
                                  ret_use_npy = batchify_fn(
                                      random_data_npy
                                  )
                                  ret_use_mx = batchify_fn(
                                      [
                                          tuple(
                                              mx.np.array(
                                                  ele[i], dtype=dtype
                                              )
                                              for i in range(
                                                  TOTAL_ELE_NUM
                                              )
                                          )
                                          for ele in random_data_npy
                                      ]
                                  )
                                  for i in range(TOTAL_ELE_NUM):
                                      if i in pad_index:
                                          assert (
                                              ret_use_npy[i][0].dtype
                                              == ret_use_mx[i][0].dtype
                                              == dtype
                                          )
                                          assert_allclose(
                                              ret_use_npy[i][
                                                  0
                                              ].asnumpy(),
                                              ret_use_mx[i][
                                                  0
                                              ].asnumpy(),
                                          )
                                      else:
                                          assert (
                                              ret_use_npy[i].dtype
                                              == ret_use_mx[i].dtype
                                              == dtype
                                          )
                                          assert_allclose(
                                              ret_use_npy[i].asnumpy(),
                                              ret_use_mx[i].asnumpy(),
                                          )
                                  for i in range(batch_size):
                                      for j in range(TOTAL_ELE_NUM):
                                          if j in pad_index:
                                              batch_data = ret_use_npy[
                                                  j
                                              ].asnumpy()
                                          else:
                                              batch_data = ret_use_npy[
                                                  j
                                              ].asnumpy()
                                          pad_length = max(
                                              ele[j].shape[axis]
                                              for ele in random_data_npy
                                          )
                                          _verify_padded_arr(
                                              batch_data[i],
                                              random_data_npy[i][j],
                                              axis,
                                              pad_val,
                                              pad_length,
                                              dtype,
                                          )
                          for _dtype in [np.float16, np.float32]:
                              shapes = [
                                  [2 for _ in range(ndim)]
                                  for _ in range(batch_size)
                              ]
                              for i in range(len(shapes)):
                                  shapes[i][axis] = np.random.randint(
                                      length_min, length_max
                                  )
                              random_data_npy = [
                                  np.random.normal(0, 1, shape).astype(
                                      dtype
                                  )
                                  for shape in shapes
                              ]
                              batchify_fn = batchify.Pad(
                                  axis=axis, val=pad_val, dtype=_dtype
                              )
                              batch_data = batchify_fn(random_data_npy)
                              batch_data_use_mx = batchify_fn(
                                  [
                                      mx.np.array(ele, dtype=dtype)
                                      for ele in random_data_npy
                                  ]
                              )
                              assert (
                                  batch_data.dtype
                                  == batch_data_use_mx.dtype
                                  == _dtype
                              )
  
file_path: gluon-nlp/tests/test_data_batchify.py
imports: import numpy
  from numpy.testing import assert_allclose
  from collections import namedtuple
  import mxnet
  from gluonnlp.data import batchify
  import pytest
label_desc: Use image node_id 6 for calling a global function with example usage: test_pad_wrap_batchify() without return types