OS:
Mac

Linux

CUDA version:
none

10.1

10.2

11.0

11.1

Command to run:

      $ pip install --upgrade jaxlib jax coax