code: def gaussian_sample(rng, mu, sigmasq):
      """Sample a diagonal Gaussian."""
      return mu + jnp.sqrt(sigmasq) * random.normal(rng, mu.shape)
  
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 2 for calling a global function with example usage: gaussian_sample(rng, mu, sigmasq) and returns: unknown