code: def model(latents):
      return VAE(latents=latents)
  
file_path: flax/examples/vae/models.py
imports: from flax import linen
  from jax import random
  import jax.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: model(latents) and returns: VAE