code: def create_model():
      """Creates a Keras model using the LeNet-5 architecture.
  
      Returns:
          model: Compiled Keras model.
      """
      # KL divergence weighted by the number of training samples, using
      # lambda function to pass as input to the kernel_divergence_fn on
      # flipout layers.
      kl_divergence_function = lambda q, p, _: tfd.kl_divergence(
          q, p
      ) / tf.cast(  # pylint: disable=g-long-lambda
          NUM_TRAIN_EXAMPLES, dtype=tf.float32
      )
  
      # Define a LeNet-5 model using three convolutional (with max pooling)
      # and two fully connected dense layers. We use the Flipout
      # Monte Carlo estimator for these layers, which enables lower variance
      # stochastic gradients than naive reparameterization.
      model = tf_keras.models.Sequential(
          [
              tfp.layers.Convolution2DFlipout(
                  6,
                  kernel_size=5,
                  padding="SAME",
                  kernel_divergence_fn=kl_divergence_function,
                  activation=tf.nn.relu,
              ),
              tf_keras.layers.MaxPooling2D(
                  pool_size=[2, 2], strides=[2, 2], padding="SAME"
              ),
              tfp.layers.Convolution2DFlipout(
                  16,
                  kernel_size=5,
                  padding="SAME",
                  kernel_divergence_fn=kl_divergence_function,
                  activation=tf.nn.relu,
              ),
              tf_keras.layers.MaxPooling2D(
                  pool_size=[2, 2], strides=[2, 2], padding="SAME"
              ),
              tfp.layers.Convolution2DFlipout(
                  120,
                  kernel_size=5,
                  padding="SAME",
                  kernel_divergence_fn=kl_divergence_function,
                  activation=tf.nn.relu,
              ),
              tf_keras.layers.Flatten(),
              tfp.layers.DenseFlipout(
                  84,
                  kernel_divergence_fn=kl_divergence_function,
                  activation=tf.nn.relu,
              ),
              tfp.layers.DenseFlipout(
                  NUM_CLASSES,
                  kernel_divergence_fn=kl_divergence_function,
                  activation=tf.nn.softmax,
              ),
          ]
      )
  
      # Model compilation.
      optimizer = tf_keras.optimizers.Adam(lr=FLAGS.learning_rate)
      # We use the categorical_crossentropy loss since the MNIST dataset contains
      # ten labels. The Keras API will then automatically add the
      # Kullback-Leibler divergence (contained on the individual layers of
      # the model), to the cross entropy loss, effectively
      # calcuating the (negated) Evidence Lower Bound Loss (ELBO)
      model.compile(
          optimizer,
          loss="categorical_crossentropy",
          metrics=["accuracy"],
          experimental_run_tf_function=False,
      )
      return model
  
imports: import os
  import warnings
  from absl import app
  from absl import flags
  import matplotlib
  from matplotlib import figure
  from matplotlib.backends import backend_agg
  import numpy
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import tf_keras
label_desc: Use image node_id 3 for calling a global function with example usage: create_model() and returns: model