add wiki sentences

working saving & loading

Merge branch 'main' of https://huggingface.co/flax-community/transformer-vae into main

fix reg loss

added rng key to the eval_step

vmapped the mmd loss function

almost working training

getting closer

add latent dim args

passing most tests

ALMOST WORKING

use latest

update todos & add steps

Merge remote-tracking branch 'origin/theo' into main

fixed the inputs for the reg_loss (still not tried it)

added the mmd loss to the loss function (not tried the whole thing yet)

closer to working

add save & eval steps

getting closer with the code

fixes

add todos

show where to add reg loss

basic code