code: def run_bnn_test(self, make_conv, make_deconv):
      # 1  Prepare Dataset
  
      train_size = 128
      batch_size = 2
      train_dataset = tf.data.Dataset.from_tensor_slices(
          tf.random.uniform(
              [train_size, 28, 28, 1], maxval=1, dtype=tf.float32
          )
      )
      train_dataset = tfn.util.tune_dataset(
          train_dataset,
          batch_size=batch_size,
          shuffle_size=int(train_size / 7),
      )
      train_iter = iter(train_dataset)
      x = next(train_iter)
      input_channels = int(x.shape[-1])
  
      # 2  Specify Model
  
      bottleneck_size = 2
  
      scale = deferred_tensor.TransformedVariable(
          1.0, softplus.Softplus()
      )
  
      bnn = tfn.Sequential(
          [
              make_conv(
                  input_channels, 32, filter_shape=5, strides=2
              ),  # [b, 14, 14, 32]
              tfn.util.flatten_rightmost(ndims=3),  # [b, 14 * 14 * 32]
              tfn.AffineVariationalReparameterization(
                  14 * 14 * 32, bottleneck_size
              ),  # [b, 2]
              lambda x: x[
                  ..., tf.newaxis, tf.newaxis, :
              ],  # [b, 1, 1, 2]
              make_deconv(
                  2, 64, filter_shape=7, strides=1, padding="valid"
              ),  # [b, 7, 7, 64]
              make_deconv(
                  64, 32, filter_shape=4, strides=4
              ),  # [2, 28, 28, 32]
              make_conv(
                  32, 1, filter_shape=2, strides=1
              ),  # [2, 28, 28, 1]
              lambda loc: independent.Independent(  # pylint: disable=g-long-lambda
                  sigmoid.Sigmoid()(normal.Normal(loc, scale)),
                  reinterpreted_batch_ndims=3,
              ),  # [b, 28, 28, 1]
          ],
          name="bayesian_autoencoder",
      )
  
      # 3  Train.
  
      def loss_fn():
          x = next(train_iter)
          nll = -tf.reduce_mean(bnn(x).log_prob(x), axis=-1)
          kl = tfn.losses.compute_extra_loss(bnn) / tf.cast(
              train_size, tf.float32
          )
          loss = nll + kl
          return loss, (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 tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import sigmoid
  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, make_deconv) and returns: loss,  nll,  kl