code: def _pad_arrs_to_max_length(
      arrs, pad_axis, pad_val, num_shards=1, use_shared_mem=False
  ):
      """Inner Implementation of the Pad batchify
      Parameters
      ----------
      arrs : list
      pad_axis : int or tuple
      pad_val : number
      num_shards : int, default 1
      use_shared_mem : bool, default False
      Returns
      -------
      ret : NDArray or a list of NDArrays
      original_length : NDArray or a list of NDArrays
      """
      if not isinstance(arrs[0], (mx.nd.NDArray, np.ndarray)):
          arrs = [np.asarray(ele) for ele in arrs]
      if isinstance(pad_axis, tuple):
          original_length = []
          for axis in pad_axis:
              original_length.append(
                  np.array([ele.shape[axis] for ele in arrs])
              )
          original_length = np.stack(original_length).T
      else:
          original_length = np.array(
              [ele.shape[pad_axis] for ele in arrs]
          )
          pad_axis = [pad_axis]
      if len(original_length) % num_shards != 0:
          logging.warning(
              "Batch size cannot be evenly split. Trying to shard %d items into %d shards",
              len(original_length),
              num_shards,
          )
      original_length = np.array_split(original_length, num_shards)
      max_lengths = [
          np.max(l, axis=0, keepdims=len(pad_axis) == 1)
          for l in original_length
      ]
      # add batch dimension
      ret_shape = [
          [
              l.shape[0],
          ]
          + list(arrs[0].shape)
          for l in original_length
      ]
      for i, shape in enumerate(ret_shape):
          for j, axis in enumerate(pad_axis):
              shape[1 + axis] = max_lengths[i][j]
      if use_shared_mem:
          ret = [
              mx.nd.full(
                  shape=tuple(shape),
                  val=pad_val,
                  ctx=mx.Context("cpu_shared", 0),
                  dtype=arrs[0].dtype,
              )
              for shape in ret_shape
          ]
          original_length = [
              mx.nd.array(
                  l, ctx=mx.Context("cpu_shared", 0), dtype=np.int32
              )
              for l in original_length
          ]
      else:
          ret = [
              mx.nd.full(
                  shape=tuple(shape), val=pad_val, dtype=arrs[0].dtype
              )
              for shape in ret_shape
          ]
          original_length = [
              mx.nd.array(l, dtype=np.int32) for l in original_length
          ]
      for i, arr in enumerate(arrs):
          if ret[i // ret[0].shape[0]].shape[1:] == arr.shape:
              ret[i // ret[0].shape[0]][i % ret[0].shape[0]] = arr
          else:
              slices = [slice(0, l) for l in arr.shape]
              ret[i // ret[0].shape[0]][i % ret[0].shape[0]][
                  tuple(slices)
              ] = arr
      if len(ret) == len(original_length) == 1:
          return ret[0], original_length[0]
      return ret, original_length
  
imports: import logging
  import numpy
  import mxnet
label_desc: Use image node_id 1 for calling a global function with example usage: _pad_arrs_to_max_length(arrs, pad_axis, pad_val, num_shards, use_shared_mem) and returns: ret,  original_length, ret,  original_length