class_docstr: Full VAE model.
code: def __call__(self, x, z_rng):
      mean, logvar = self.encoder(x)
      z = reparameterize(z_rng, mean, logvar)
      recon_x = self.decoder(z)
      return recon_x, mean, logvar
  
file_path: flax/examples/vae/models.py
imports: from flax import linen
  from jax import random
  import jax.numpy
label_desc: Use image node_id 2 for calling the VAE obj's underlying member method code with example usage: obj.__call__(x, z_rng) and returns: recon_x,  mean,  logvar