code: def sample_diag_gaussian(mean, log_std, rs):
      return rs.randn(*mean.shape) * np.exp(log_std) + mean
  
imports: from __future__ import absolute_import, division
  from __future__ import print_function
  import autograd.numpy
  import autograd.numpy.random
  import autograd.scipy.stats.norm
  from autograd.scipy.special import expit
  from autograd import grad
  from autograd.misc.optimizers import adam
  from data import load_mnist, save_images
label_desc: Use image node_id 3 for calling a global function with example usage: sample_diag_gaussian(mean, log_std, rs) and returns: unknown