code: def eval_f(params, images, z, z_rng, latents):
      def eval_model(vae):
          recon_images, mean, logvar = vae(images, z_rng)
          comparison = jnp.concatenate(
              [
                  images[:8].reshape(-1, 28, 28, 1),
                  recon_images[:8].reshape(-1, 28, 28, 1),
              ]
          )
  
          generate_images = vae.generate(z)
          generate_images = generate_images.reshape(-1, 28, 28, 1)
          metrics = compute_metrics(recon_images, images, mean, logvar)
          return metrics, comparison, generate_images
  
      return nn.apply(eval_model, models.model(latents))(
          {"params": params}
      )
  
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 5 for calling a global function with example usage: eval_f(params, images, z, z_rng, latents) and returns: nn, metrics,  comparison,  generate_images