def chunk(array, num_memory_steps):
    """Chunk an array into size of N steps.

    The array of size N x k_1 x ... k_n will be reshaped to be of size
    Batch x num_memory_steps x k_1 x ... k_n.

    Parameters
    ----------
    array: Array.
        Array to reshape.
    num_memory_steps: int.
        Number of steps to chunk the batch.

    Returns
    -------
    array: Array.
        Chunked Array.
    """
    batch_size = array.shape[0] // num_memory_steps
    return array.reshape(batch_size, num_memory_steps, *array.shape[1:])