code: def diag_gaussian_sample(rng, mean, log_std):
      # Take a single sample from a diagonal multivariate Gaussian.
      return mean + jnp.exp(log_std) * random.normal(rng, mean.shape)
  
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 1 for calling a global function with example usage: diag_gaussian_sample(rng, mean, log_std) and returns: unknown