🐍
google/jax を Poetry でインストールする
背景
google/jax が PEP 503 に対応していなくて困った
Poetry 側にも issue がある
どうやら pip の --find-links 相当のものが 1.2.0 から実装されたらしい
このバージョン以降ならインストールできそう
インストール方法
$ poetry source add jax https://storage.googleapis.com/jax-releases/jax_releases.html
$ poetry add --source jax jaxlib
$ poetry add jax
これで 0.3.20+cuda11.cudnn82 の jaxlib がインストールされて、jax から cuda を認識できるようになった
仕組み
jax[cuda] や jax[tpu] は extras という仕組みを使ったインストール方法
パッケージ側では setup.py に extras_require を書くことで指定できる
jax の仕組みは、extras に応じて jaxlib の実装を切り替えることによってバックエンドを使い分けているっぽい
Discussion