CPUによる計算環境において，JAXの行列計算速度がnumpyに劣ってしまう問題

Question

前提・実現したいこと

私は機械学習のフレームであるjaxを学び始めております．
そこで，以下のkaggleにあるjaxの紹介を基に勉強しておりました．
https://www.kaggle.com/aakashnain/tf-jax-tutorials-part-4-jax-and-devicearray/notebook

jaxの強みは高速な自動微分であると思っていたのですが，
numpyよりも遅くなってしまう事例が発生しました．
皆様には，なぜjaxの計算が遅くなってしまったのかについて，
ご意見をお聞きしたいと思っております．
（追記　質問内容が意図と異なっていたため修正いたします）
google colab上でCPUによる計算を行わせたところ，jaxの行列計算速度がnumpyよりも遅くなってしまう事例が発生しました．
jaxの強みは高速な自動微分であると思っていたため，CPU上では計算速度が低くなってしまうことを不思議に思い，
また，jaxの知識が浅くよくわかっていないため，皆様へご意見をお聞きしたいと思います．
jaxはCPUの計算環境においては，計算速度がnumpyよりも劣ってしまうのでしょうか？

発生している問題・エラーメッセージ

https://www.kaggle.com/aakashnain/tf-jax-tutorials-part-4-jax-and-devicearray/notebook
の[12]における

# Dot product on ndarray
start_time = time.time()
res = np.dot(array1, array1)
print(f"Time taken by dot product op on ndarrays: {time.time()-start_time:.2f} seconds")

# Dot product on DeviceArray
start_time = time.time()
res = jnp.dot(array2, array2)
print(f"Time taken by dot product op on DeviceArrays: {time.time()-start_time:.2f} seconds")

は，ページ内では
Time taken by dot product op on ndarrays: 7.94 seconds
Time taken by dot product op on DeviceArrays: 0.02 seconds
とDeviceArrayの高速性が発揮していますが，

私の環境（Google Colab/ CPU）で同等のコードで実行したところ，
Time taken by dot product op on ndarrays: 14.60 seconds
Time taken by dot product op on DeviceArrays: 16.78 seconds
と，numpyよりも計算が遅い結果となってしまいました．

試したこと

私のコードは，全てからコピペしたものであり，全く同じコードです．
それなのに，計算速度に違いが生まれてしまうことが不思議でなりません．
唯一出たエラーコードは
[1]を実行したとき：/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py:8: UserWarning: Config option use_jedi not recognized by IPCompleter.
[2]を実行したとき：WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
jaxはCPUの計算環境においては，計算速度がnumpyよりも劣ってしまうのでしょうか？
お詳しい方がいらっしゃいましたら，ご意見をいただけますと幸いです．
よろしくお願いいたします．

Comments

[quickquip]

> WARNING:absl:No GPU/TPU found, falling back to CPU.
ここを読んでますか?

[guratan]

ご返信ありがとうございます．
そのエラーコードは，
「GPUが使える状態ではないから，CPUで計算します」という理解のみでした．
CPUで計算を行ったとしても，Jaxの計算速度が高速である，という考えだったため，
この質問を行った次第です．
GPUが利用できる環境であれば，jaxの計算速度が本来のものになる，ということでしょうか？
知識が浅く，申し訳ございません．

[jbpb0]

参考
https://stackoverflow.com/questions/63672151/why-does-jax-numpy-dot-run-slower-than-numpy-dot-on-cpu

[quickquip]

> CPUで計算を行ったとしても，Jaxの計算速度が高速である，という考えだったため，
質問の「jaxは計算の違いで」のあたりがそういう意図なのですね。
その部分を（ppaulさんの回答と食い違わない範囲で）書き換えた方がいいかもと思います。今からでも。

[guratan]

>その部分を（ppaulさんの回答と食い違わない範囲で）書き換えた方がいいかもと思います。今からでも。
おっしゃる通りです．私の質問文の説明が不足しておりました...
そのように修正しようと思います．ご指摘ありがとうございます．

Answer 1

GPUもTPUも付いていないハードで実行するなら、JAXがインテルのMKLを使ったnumpyよりも遅いことはいくらでもあるでしょう。

JAX入門～高速なNumPyとして使いこなすためのチュートリアル～などには明確には書かれていませんが、結局jaxはGoogleが開発したのために開発されたライブラリで、GPUにも使えますという位置づけのように見えます。
インテルは自社のハード向けにMKLに資金を投入していますから、JAXでCPU向けにJITしたのではMKLとの勝負ではそれほど強いとは思えませんね。

JAXを使いたければればゲーミングPCに買い換えるか、Google Colaｂを使うべきでしょう。

Comments

[guratan]

ご返信ありがとうございます．
やはり，CPUではjaxの性能を発揮できないようですね...
大変参考になりました．

[jbpb0]

> CPUではjaxの性能を発揮できない

https://qiita.com/koshian2/items/44a871386576b4f80aff
の「パフォーマンス比較～CPU～」では、大きなサイズの配列の要素積では、jitありJAXの方がNumPyよりも高速だったと書かれてます
CPUでも、JAXが優位になるような条件の場合は、使う価値はあると思いますよ

Answer 2

google colab pro + GPU使用に変えたところ，
Time taken by dot product op on ndarrays: 14.93 seconds
Time taken by dot product op on DeviceArrays: 3.53 seconds
とjaxの計算速度が高速になりました．
jaxを使用するときは，機械学習なら当然ですがGPU/TPUを使用することが前提のようですね．

ご返信いただいた皆様，ありがとうございました．