Deep Learning with JAX and Flax/

...

/

Installing Packages

Installing Packages

Learn how to install the packages on your local machine.

We'll cover the following...

Installing JAX

JAX can be installed from the Python Package Index (PIP) as follows:

pip install jax
Installing the jax framework

JAX is preinstalled on Google Colaboratory.

Setting up TPUs on Google Colaboratory

You’ll need to set up JAX to use TPUs on Colab. This is done by executing the following code. Ensure that you have changed the runtime to TPU by going to “Runtime” -> “Change ...