5.5 Implementing BBB with TensorFlow

In this section, we’ll see how to implement BBB in TensorFlow. Some of the code you’ll see will be familiar; the core concepts of layers, loss functions, and optimizers will be very similar to what we covered in Chapter 3, Fundamentals of Deep Learning. Unlike the examples in Chapter 3, Fundamentals of Deep Learning, we’ll see how we can create neural networks capable of probabilistic inference.

Step 1: Importing packages

We start by importing the relevant packages. Importantly, we will import tensorflow-probability, which will provide us with the layers of the network that replace the point-estimate with a distribution and implement the reparameterization trick. We also set the global parameter for the number of inferences, which will determine how often we sample from the network later:

 
import tensorflow as tf  
import numpy as np  
import matplotlib.pyplot as plt  
import...