code: def main(_):
      flags.FLAGS.alsologtostderr = True
      config = Config()
  
      @hk.transform
      def model(x):
          vae = VariationalAutoEncoder(
              encoder=Encoder(latent_size=10),
              decoder=Decoder(output_shape=x.shape[1:]),
          )
          return vae(x)
  
      @jax.jit
      def loss_fn(params, rng_key, batch: Batch) -> jax.Array:
          """ELBO loss: E_p[log(x)] - KL(d||q), where p ~ Be(0.5) and q ~ N(0,1)."""
  
          # Run the model on the inputs.
          _, mean, var, logits = model.apply(
              params, rng_key, batch.image
          )
  
          # Bernoulli log-likelihood (assumes `image` is binarised).
          log_likelihood = jnp.einsum(
              "b...->b",
              batch.image * logits - jnp.logaddexp(0.0, logits),
          )
  
          # KL divergence between Gaussians N(mean, std) and N(0, 1).
          kl = 0.5 * jnp.sum(
              -jnp.log(var) - 1.0 + var + jnp.square(mean), axis=-1
          )
  
          # Loss is the negative evidence lower-bound.
          return -jnp.mean(log_likelihood - kl)
  
      optimizer = optax.adam(config.learning_rate)
  
      @jax.jit
      def update(state: TrainingState, batch: Batch) -> TrainingState:
          """Performs a single SGD step."""
          rng_key, next_rng_key = jax.random.split(state.rng_key)
          gradients = jax.grad(loss_fn)(state.params, rng_key, batch)
          updates, new_opt_state = optimizer.update(
              gradients, state.opt_state
          )
          new_params = optax.apply_updates(state.params, updates)
          return TrainingState(new_params, new_opt_state, next_rng_key)
  
      # Load datasets.
      train_dataset = load_dataset(
          "train", config.batch_size, config.seed
      )
      eval_datasets = {
          "train": load_dataset(
              "train", config.batch_size, config.seed
          ),
          "valid": load_dataset(
              "validation", config.batch_size, config.seed
          ),
      }
  
      # Initialise the training state.
      initial_rng_key = jax.random.PRNGKey(config.seed)
      initial_params = model.init(
          initial_rng_key, next(train_dataset).image
      )
      initial_opt_state = optimizer.init(initial_params)
      state = TrainingState(
          initial_params, initial_opt_state, initial_rng_key
      )
  
      # Run training and evaluation.
      for step in range(config.training_steps):
          state = update(state, next(train_dataset))
  
          if step % config.eval_every == 0:
              for split, ds in eval_datasets.items():
                  loss = loss_fn(state.params, state.rng_key, next(ds))
                  logging.info(
                      {
                          "step": step,
                          "split": split,
                          "elbo": -jax.device_get(loss).item(),
                      }
                  )
  
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 2 for calling a global function with example usage: main(_) and returns: vae, unknown, TrainingState