Jax sounds like an impressive project, thanks for working on it. That said: on Ubuntu 18.04, this happens ➜ python Python 3.6.9 (default, Oct 8 2020, 12:12:24) [GCC 8.4.0] on linux Type "help", "co...

Jax doesn't see my GPU, even though Pytorch does · Issue #5231 · jax-ml/jax