code: def batch_elbo(logprob, rng, params, num_samples):
      # Average over a batch of random samples.
      rngs = random.split(rng, num_samples)
      vectorized_elbo = vmap(
          partial(elbo, logprob), in_axes=(0, None, None)
      )
      return jnp.mean(vectorized_elbo(rngs, *params))
  
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 4 for calling a global function with example usage: batch_elbo(logprob, rng, params, num_samples) and returns: jnp