* add discrete VAE

* also return z_mode for discrete vae

* add ST gumbel-softmax

* update

* add scheduler option for the temperature

* address comments