code: def create_model(num_samples, num_dimensions):
      """Creates a Keras model for logistic regression.
  
      Args:
       num_samples: Integer for number of training samples.
       num_dimensions: Integer for number of features in dataset.
  
      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_samples, dtype=tf.float32
      )
  
      # Define a logistic regression model as a Bernoulli distribution
      # parameterized by logits from a single linear layer. We use the Flipout
      # Monte Carlo estimator for the layer: this enables lower variance
      # stochastic gradients than naive reparameterization.
      input_layer = tf_keras.layers.Input(shape=num_dimensions)
      dense_layer = tfp.layers.DenseFlipout(
          units=1,
          activation="sigmoid",
          kernel_posterior_fn=tfp.layers.default_mean_field_normal_fn(),
          bias_posterior_fn=tfp.layers.default_mean_field_normal_fn(),
          kernel_divergence_fn=kl_divergence_function,
      )(input_layer)
  
      # Model compilation.
      model = tf_keras.Model(inputs=input_layer, outputs=dense_layer)
      optimizer = tf_keras.optimizers.Adam(lr=FLAGS.learning_rate)
      # We use the binary_crossentropy loss since this toy example contains
      # two 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="binary_crossentropy", metrics=["accuracy"]
      )
      return model
  
file_path: probability/tensorflow_probability/examples/logistic_regression.py
imports: import os
  from absl import app
  from absl import flags
  from matplotlib import cm
  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(num_samples, num_dimensions) and returns: model