Installing Packages
Learn how to install the packages on your local machine.
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 ...