flax

6df99e44 - Merge branch 'pixelcnn' of https://github.com/marcvanzee/flax into pixelcnn