code: def elbo(logprob, rng, mean, log_std):
      # Single-sample Monte Carlo estimate of the variational lower bound.
      sample = diag_gaussian_sample(rng, mean, log_std)
      return logprob(sample) - diag_gaussian_logpdf(
          sample, mean, log_std
      )
  
file_path: jax/examples/advi.py
imports: from functools import partial
  import matplotlib.pyplot
  from jax import jit, grad, vmap
  from jax import random
  from jax.example_libraries import optimizers
  import jax.numpy
  import jax.scipy.stats.norm
label_desc: Use image node_id 3 for calling a global function with example usage: elbo(logprob, rng, mean, log_std) and returns: unknown