Fix failing float32 tests
Make 'learn_prior' the responsibility of VAE
Fix bad default 'mu' layer for DiagonalGaussian