


act_fn (str, optional, defaults to silu) —
Activation function


latent_channels (int, optional, defaults to 4) —
Latent space channels


norm_num_groups (int, optional, defaults to 32) —
Norm num group


sample_size (int, optional, defaults to 32) —
Sample input size


dtype (jnp.dtype, optional, defaults to jnp.float32) —
parameters dtype



Flax Implementation of Variational Autoencoder (VAE) model with KL loss from the paper Auto-Encodin