code: def image_sample(rng, params, nrow, ncol):
      """Sample images from the generative model."""
      _, dec_params = params
      code_rng, img_rng = random.split(rng)
      logits = decode(
          dec_params, random.normal(code_rng, (nrow * ncol, 10))
      )
      sampled_images = random.bernoulli(
          img_rng, jnp.logaddexp(0.0, logits)
      )
      return image_grid(nrow, ncol, sampled_images, (28, 28))
  
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 5 for calling a global function with example usage: image_sample(rng, params, nrow, ncol) and returns: image_grid