Bug fixes for LARS:
* When used as proposal of VAE, correctly use current sample.
* Exp moving average was double exponentiating.
Switch to TFDS for all datasets and compute mean on the fly.
Implement some Conv networks and versions of the models.
Remove flat datasets. Now all methods should handle the flattening themselves.
Fix bugs with cacheing in dataset.py.
Bug fixes and updates for TF for small_problems.py.

PiperOrigin-RevId: 275357997