JAX: Difference between revisions
Remove out of date note, add link to nixpkgs and other wiki pages |
m jax.lib.xla_bridge module will be removed in JAX v0.9.0; all its APIs were deprecated and removed by JAX v0.8.0 |
||
| Line 41: | Line 41: | ||
You can test that JAX is using the GPU as intended with | You can test that JAX is using the GPU as intended with | ||
# after version 0.8.0 | |||
python -c "import jax.extend as ex; print(ex.backend.get_backend().platform)" | |||
# before version 0.8.0 | |||
python -c "from jax.lib import xla_bridge; print(xla_bridge.get_backend().platform)" | python -c "from jax.lib import xla_bridge; print(xla_bridge.get_backend().platform)" | ||