はじめに

コードを触っている中で、動作はできているのになぜか処理が重いという現象がありました。
その原因はまた別だったのですが、副産物としてタイトル通りの問題を見つけたので共有しておきます。

尚、ここではGeneratorは何かは説明しません。詳しくは以下の記事を参照お願いします。

確認

試しに100000x2の配列を100回シャッフルしてみて、合計と平均時間を見てみます。
ざっと書くとこんな感じ。

import random
import time
random.seed(0)

import numpy as np
np.random.seed(0)

rng = np.random.default_rng(0)


def shuffle(func, array, count):
    total = 0.
    for _ in range(count):
        start = time.time()
        func.shuffle(array)
        end = time.time()
        total += end - start
    print('\tTotal time: {:.6f} seconds'.format(total))
    print('\tAverage time: {:.6f} seconds'.format(total / count))


base_array = rng.uniform(size=(100000, 2))

print('np.random.shuffle')
shuffle(np.random, base_array.copy(), 100)

print('rng.shuffle')
shuffle(rng, base_array.copy(), 100)

1次元

1次元においてはむしろGeneratorの方が早いです。

np.random.shuffle
        Total time: 0.344189 seconds
        Average time: 0.003442 seconds
rng.shuffle
        Total time: 0.312505 seconds
        Average time: 0.003125 seconds

2次元

2次元の結果(100000x2)は以下の通りになりました。

np.random.shuffle
        Total time: 25.462759 seconds
        Average time: 0.254628 seconds
rng.shuffle
        Total time: 29.436836 seconds
        Average time: 0.294368 seconds

だいたい0.04秒遅れてる感じになるでしょうか。次に1000000x2でやってみます

np.random.shuffle
        Total time: 286.438086 seconds
        Average time: 2.864381 seconds
rng.shuffle
        Total time: 366.751766 seconds
        Average time: 3.667518 seconds

こちらは0.8秒ほどでしょうか。けっこう差が開いてきますね。

原因を探る

さて、numpyの実装を見てみます。まずはnp.random.shuffleから。

            #...
            buf = np.empty_like(x[0, ...])
            with self.lock:
                for i in reversed(range(1, n)):
                    j = random_interval(&self._bitgen, i)
                    if i == j:
                        continue  # i == j is not needed and memcpy is undefined.
                    buf[...] = x[j]
                    x[j] = x[i]
                    x[i] = buf
            #...

次にnp.random.Generator.shuffleを。

            #...
            with self.lock:
                for i in reversed(range(1, len(x))):
                    j = random_interval(&self._bitgen, i)
                    if i == j:
                        # i == j is not needed and memcpy is undefined.
                        continue
                    buf[...] = x[j, ...]
                    x[j, ...] = x[i, ...]
                    x[i, ...] = buf
            #...

さて、両者ともcで実装されたrandom_intervalを呼び出しています。まぁ、これは同じものを呼び出しているので、差はないでしょう。
一方で気になるのがcontinue以降です。bufとxのswapをしていますが、Generator側はelipsisをすべてに使っています。一方でnp.randomはelipsisは最初のbufの代入以外使っていません。
numpyの配列のelipsisは次元の省略の指定になりますが、これが原因だとすれば、array側にこの指定処理が入ってしまって遅くなってる、という感じになるのでしょうか。

おわりに

shuffleは基本的に1回ぐらいしか使わない or 1次元の場合がほとんどですが、例えば機械学習のstep中に大量の2次元配列の座標をシャッフルしてる感じとかだと、こういうところは遅くなる原因の1つになるかとおもいます(だいぶレアなケースですが…)。
ソースコード見る限りelipsisの指定の有無が原因っぽくは見えますが、どうなんでしょう。アルゴリズムは同じだと思いますが…