fixes #1

Author: activatedgeek

README.md

@@ -1,18 +1,23 @@
 # SWAG in Optax
 
+[![PyPI version](https://badge.fury.io/py/optax-swag.svg)](https://pypi.org/project/optax-swag/)
+
 This package implements [SWAG](https://arxiv.org/abs/1902.02476) 
 as an [Optax](https://optax.readthedocs.io/) transform to allow
 usage with [JAX](https://jax.readthedocs.io/).
 
 ## Installation
 
-For now, the only available mode of installation is directly from source as
+Install from `pip` as:
+```shell
+pip install optax-swag
 ```
+
+To install the latest directly from source, run
+```shell
 pip install git+https://github.com/activatedgeek/optax-swag.git
 ```
 
-**TODO**: A PyPI package will be available soon.
-
 ## Usage
 
 To start updating the iterate statistics, use [chaining](https://optax.readthedocs.io/en/latest/api.html#chain) as