def _pad_arrs_to_max_length(arrs, pad_axis, pad_val, use_shared_mem, dtype):
if isinstance(arrs[<NUM_LIT:0>], mx.nd.NDArray):<EOL><INDENT>dtype = arrs[<NUM_LIT:0>].dtype if dtype is None else dtype<EOL>arrs = [arr.asnumpy() for arr in arrs]<EOL><DEDENT>elif not isinstance(arrs[<NUM_LIT:0>], np.ndarray):<EOL><INDENT>arrs = [np.asarray(ele) for ele in arrs]<EOL><DEDENT>else:<EOL><INDENT>dtype = arrs[<NUM_LIT:0>].dtype if dtype is None else dtype<EOL><DEDENT>original_length = [ele.shape[pad_axis] for ele in arrs]<EOL>max_size = max(original_length)<EOL>ret_shape = list(arrs[<NUM_LIT:0>].shape)<EOL>ret_shape[pad_axis] = max_size<EOL>ret_shape = (len(arrs), ) + tuple(ret_shape)<EOL>ret = np.full(shape=ret_shape, fill_value=pad_val, dtype=dtype)<EOL>for i, arr in enumerate(arrs):<EOL><INDENT>if arr.shape[pad_axis] == max_size:<EOL><INDENT>ret[i] = arr<EOL><DEDENT>else:<EOL><INDENT>slices = [slice(None) for _ in range(arr.ndim)]<EOL>slices[pad_axis] = slice(<NUM_LIT:0>, arr.shape[pad_axis])<EOL>if slices[pad_axis].start != slices[pad_axis].stop:<EOL><INDENT>slices = [slice(i, i + <NUM_LIT:1>)] + slices<EOL>ret[tuple(slices)] = arr<EOL><DEDENT><DEDENT><DEDENT>ctx = mx.Context('<STR_LIT>', <NUM_LIT:0>) if use_shared_mem else mx.cpu()<EOL>ret = mx.nd.array(ret, ctx=ctx, dtype=dtype)<EOL>original_length = mx.nd.array(original_length, ctx=ctx, dtype=np.int32)<EOL>return ret, original_length<EOL>
Inner Implementation of the Pad batchify

    Parameters
    ----------
    arrs : list
    pad_axis : int
    pad_val : number
    use_shared_mem : bool, default False

    Returns
    -------
    ret : NDArray
    original_length : NDArray