Fix VAE training stability

[0xtristan]

Currently @sholtodouglas uses dynamic gradient clipping (gradient clipped if 4x the previous norm). This still seems to cause some issues with loss spikes.

Other options are static clipping, gradient penalties, weight regularisation/normalisation/decay

A personal favourite of mine that worked very well (and is also leveraged in NVAE) is SpectralNorm. TFA has an implementation here:
https://www.tensorflow.org/addons/api_docs/python/tfa/layers/SpectralNormalization