機械学習ライブラリJAXのインストールで詰まった話

JAXという機械学習ライブラリがpythonから使えると聞いたのでインストールしました
インストール時に詰まったところのメモです

pipを使ってJAXをインストールしようとしました

pip install jax jaxlib

最初にではエラーはこちら：

ValueError: jaxlib is version 0.1.23, but this version of jax requires version 0.1.37.

これはエラーメッセージの下の方に解決方法が書いてあるので，素直に実行します

pip install --upgrade pip
pip install --upgrade jax jaxlib

次に出たエラーはこちら：

AttributeError: module 'numpy' has no attribute 'histogram_bin_edges'

エラーメッセージからnumpyのエラーらしいので，numpyのupgradeをします

pip install --upgrade numpy

最後に出たエラーはこちら：

ModuleNotFoundError: No module named 'numpy.testing.decorators'

エラーメッセージの上の方を見るとnumpyではなくscipyのエラーらしいので，scipyのupgradeをします

pip install --upgrade scipy

これで動きました