I added a notebook that performs inference on a VQGAN using Suraj
Patil's Jax implementation. It demonstrates the use on a ~8K image
dataset, but it's still slow.

To Do:
- Performance improvements. Data loading or excessive TPU to CPU
transfer might be a problem. I think resizing can be performed in JAX,
so that's something to test.
- Save results to disk, possibly incrementally.
- Make sure we process all images even if the last batch is incomplete
(we are ignoring it if that's the case).