code: def train_and_evaluate(config: ml_collections.ConfigDict):
      """Train and evaulate pipeline."""
      rng = random.key(0)
      rng, key = random.split(rng)
  
      ds_builder = tfds.builder("binarized_mnist")
      ds_builder.download_and_prepare()
  
      logging.info("Initializing dataset.")
      train_ds = input_pipeline.build_train_set(
          config.batch_size, ds_builder
      )
      test_ds = input_pipeline.build_test_set(ds_builder)
  
      logging.info("Initializing model.")
      init_data = jnp.ones((config.batch_size, 784), jnp.float32)
      params = models.model(config.latents).init(key, init_data, rng)[
          "params"
      ]
  
      state = train_state.TrainState.create(
          apply_fn=models.model(config.latents).apply,
          params=params,
          tx=optax.adam(config.learning_rate),
      )
  
      rng, z_key, eval_rng = random.split(rng, 3)
      z = random.normal(z_key, (64, config.latents))
  
      steps_per_epoch = (
          ds_builder.info.splits["train"].num_examples
          // config.batch_size
      )
  
      for epoch in range(config.num_epochs):
          for _ in range(steps_per_epoch):
              batch = next(train_ds)
              rng, key = random.split(rng)
              state = train_step(state, batch, key, config.latents)
  
          metrics, comparison, sample = eval_f(
              state.params, test_ds, z, eval_rng, config.latents
          )
          vae_utils.save_image(
              comparison, f"results/reconstruction_{epoch}.png", nrow=8
          )
          vae_utils.save_image(
              sample, f"results/sample_{epoch}.png", nrow=8
          )
  
          print(
              "eval epoch: {}, loss: {:.4f}, BCE: {:.4f}, KLD: {:.4f}".format(
                  epoch + 1,
                  metrics["loss"],
                  metrics["bce"],
                  metrics["kld"],
              )
          )
  
file_path: flax/examples/vae/train.py
imports: from absl import logging
  from flax import linen
  import input_pipeline
  import models
  import utils
  from flax.training import train_state
  import jax
  from jax import random
  import jax.numpy
  import ml_collections
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 6 for calling a global function with example usage: train_and_evaluate(config) without return types