code: def gaussian_kl(mu, sigmasq):
      """KL divergence from a diagonal Gaussian to the standard Gaussian."""
      return -0.5 * jnp.sum(
          1.0 + jnp.log(sigmasq) - mu**2.0 - sigmasq
      )
  
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 1 for calling a global function with example usage: gaussian_kl(mu, sigmasq) and returns: unknown