code: def run_bnn_test(self, make_conv):
      # 1  Prepare Dataset
  
      train_size = 128
      batch_size = 2
      evidence_shape = [28, 28, 1]
      target_shape = [10]
  
      train_dataset = tf.data.Dataset.from_tensor_slices(
          (
              tf.random.uniform(
                  [train_size] + evidence_shape,
                  maxval=1,
                  dtype=tf.float32,
              ),
              tf.math.softmax(
                  tf.random.normal([train_size] + target_shape)
              ),
          )
      )
      train_dataset = tfn.util.tune_dataset(
          train_dataset,
          batch_size=batch_size,
          shuffle_size=int(train_size / 7),
      )
  
      # 2  Specify Model
  
      scale = deferred_tensor.TransformedVariable(
          1.0, softplus.Softplus()
      )
      bnn = tfn.Sequential(
          [
              make_conv(
                  evidence_shape[-1], 32, filter_shape=7, strides=2
              ),  # [b, 14, 14, 32]
              tfn.util.flatten_rightmost(ndims=3),  # [b, 14 * 14 * 32]
              tfn.AffineVariationalReparameterization(
                  14 * 14 * 32, np.prod(target_shape) - 1
              ),  # [b, 9]
              lambda loc: softmax_centered.SoftmaxCentered()(  # pylint: disable=g-long-lambda
                  independent.Independent(
                      normal.Normal(loc, scale),
                      reinterpreted_batch_ndims=1,
                  )
              ),  # [b, 10]
          ],
          name="bayesian_autoencoder",
      )
  
      self.evaluate([v.initializer for v in bnn.trainable_variables])
  
      # 3  Train.
      n = tf.cast(train_size, tf.float32)
      train_iter = iter(train_dataset)
  
      def loss_fn():
          x, y = next(train_iter)
          nll = -tf.reduce_mean(bnn(x).log_prob(y), axis=-1)
          kl = tfn.losses.compute_extra_loss(bnn) / n
          return nll + kl, (nll, kl)
  
      opt = tf_keras.optimizers.Adam()
      fit_op = tfn.util.make_fit_op(
          loss_fn, opt, bnn.trainable_variables
      )
      for _ in range(2):
          loss, (nll, kl) = fit_op()
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softmax_centered
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.experimental import nn
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.util import deferred_tensor
label_desc: Use image node_id 1 for calling the BnnEndToEnd obj's underlying member method code with example usage: obj.run_bnn_test(make_conv) and returns: unknown,  nll,  kl