code: def elbo(rng, params, images):
      """Monte Carlo estimate of the negative evidence lower bound."""
      enc_params, dec_params = params
      mu_z, sigmasq_z = encode(enc_params, images)
      logits_x = decode(
          dec_params, gaussian_sample(rng, mu_z, sigmasq_z)
      )
      return bernoulli_logpdf(logits_x, images) - gaussian_kl(
          mu_z, sigmasq_z
      )
  
file_path: jax/examples/mnist_vae.py
imports: import os
  import time
  import matplotlib.pyplot
  import jax
  import jax.numpy
  from jax import jit, grad, lax, random
  from jax.example_libraries import optimizers
  from jax.example_libraries import stax
  from jax.example_libraries.stax import Dense, FanOut, Relu, Softplus
  from examples import datasets
label_desc: Use image node_id 4 for calling a global function with example usage: elbo(rng, params, images) and returns: unknown