Updating readme with quick install instructions

PiperOrigin-RevId: 344981033

README.md

@@ -1,6 +1,7 @@
 # Flax: A neural network ecosystem for JAX designed for flexibility
 
 [**Overview**](#overview)
+| [**Quick install**](#quick-install)
 | [**What does Flax look like?**](#what-does-flax-look-like)
 | [**Documentation**](https://flax.readthedocs.io/)
 
@@ -41,6 +42,29 @@ comes with everything you need to start your research, including:
 
 * **Fast, tuned large-scale end-to-end examples**: CIFAR10, ResNet on ImageNet, Transformer LM1b
 
+## Quick install
+
+You will need Python 3.6 or later and a working [JAX](https://github.com/google/jax/blob/master/README.md)
+installation (with or without GPU support, see instructions there). For a
+CPU-only version:
+
+```
+> pip install --upgrade pip # To support manylinux2010 wheels.
+> pip install --upgrade jax jaxlib # CPU-only
+```
+
+Then install Flax from PyPi:
+
+```
+> pip install flax
+```
+
+To upgrade to the latest version of Flax, you can use:
+
+```
+> pip install --upgrade git+https://github.com/google/flax.git
+```
+
 ## What does Flax look like?
 
 We provide here two examples using the Flax API: a simple multi-layer perceptron and a CNN. To learn more about the `Module` abstraction, please check our [docs](https://flax.readthedocs.io/).