Use tfds for Cifar10 in VQ-VAE example, default to EMA and include cell outputs.
Use tfds for Cifar10 in VQ-VAE example, default to EMA and include cell outputs.

PiperOrigin-RevId: 311517163
Change-Id: Ie8ebc0f95c14d6e70118f48695be65ad7773a13b

deepmind/sonnet,deepmind/sonnet