code: def diag_gaussian_logpdf(x, mean, log_std):
      # Evaluate a single point on a diagonal multivariate Gaussian.
      return jnp.sum(vmap(norm.logpdf)(x, mean, jnp.exp(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 2 for calling a global function with example usage: diag_gaussian_logpdf(x, mean, log_std) and returns: jnp