class_docstr: Main VAE model class.
docstr: Forward pass of the variational autoencoder.
docstr_tok: ["Forward","pass","of","the","variational","autoencoder","."]
code: def __call__(self, x: jax.Array) -> VAEOutput:
      """Forward pass of the variational autoencoder."""
      x = x.astype(jnp.float32)
      mean, stddev = self.encoder(x)
      z = mean + stddev * jax.random.normal(
          hk.next_rng_key(), mean.shape
      )
      logits = self.decoder(z)
  
      p = jax.nn.sigmoid(logits)
      image = jax.random.bernoulli(hk.next_rng_key(), p)
  
      return VAEOutput(image, mean, jnp.square(stddev), logits)
  
imports: from collections.abc import Iterator, Sequence
  import dataclasses
  from typing import NamedTuple
  from absl import app
  from absl import flags
  from absl import logging
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling the VariationalAutoEncoder obj's underlying member method code with example usage: obj.__call__(x) and returns: VAEOutput