code: def compute_metrics(recon_x, x, mean, logvar):
      bce_loss = binary_cross_entropy_with_logits(recon_x, x).mean()
      kld_loss = kl_divergence(mean, logvar).mean()
      return {
          "bce": bce_loss,
          "kld": kld_loss,
          "loss": bce_loss + kld_loss,
      }
  
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 3 for calling a global function with example usage: compute_metrics(recon_x, x, mean, logvar) and returns: dict