
```
$ poetry add jax jaxlib
```

とした後、以下のコードで jax で cpu と gpu どちらを使っているかを確認できる。

```bash
$ python
>>> from jax.lib import xla_bridge
>>> print(xla_bridge.get_backend().platform)
>>> cpu
>>> これでは gpu を認識していない。
```

以下のように書いてみてもだめだった。jax が`>0.3.25,<0.4.0`と合わないというエラー。

かなりググったが、poetry では無理そうだった。

```bash
$ poetry add "jax[cuda11_cudnn82]"
Using version ^0.3.25 for jax

Updating dependencies
Resolving dependencies... (0.0s)

Because no versions of jax match >0.3.25,<0.4.0
and jax (0.3.25) depends on jaxlib (0.3.25+cuda11.cudnn82), jax (>=0.3.25,<0.4.0) requires jaxlib (0.3.25+cuda11.cudnn82).
So, because no versions of jaxlib match 0.3.25+cuda11.cudnn82
and ml-library depends on jax (^0.3.25), version solving failed.
```

以下のように pip でインストールしたら gpu を認識してくれた。

```
$ pip install "jax[cuda11_cudnn82]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
$ python
>>> from jax.lib import xla_bridge
>>> print(xla_bridge.get_backend().platform)
>>> cpu
```


PoetryでJAXのGPU版をインストールできないときの解決法