After equation 2 in the paper, the author writes:

> where $\mathbf{w}^{(i)}$ denotes the $i^{th}$ Monte Carlo sample drawn from the **variational posterior** $q(\mathbf{w}^{(i)}|\theta)$.

In section 3.2 titled "Gaussian variational posterior":

> Suppose that the **variational posterior is a diagonal Gaussian distribution,  then a sample of the weights $\mathbf{w}$ can be obtained by sampling a unit Gaussian, shifting it by a mean $\mu$ and scaling by a standard deviation $\sigma$**. We parameterise the standard deviation pointwise as $\sigma = \log(1 + \exp(\rho))$ and so $\sigma$ is always non-negative. The variational posterior parameters are $\theta = (\mu,\rho)$. Thus the transform from a sample of parameter-free noise and the variational posterior parameters that yields a posterior sample of the weights $\mathbf{w}$ is: $\mathbf{w} = t(\theta,\epsilon) = \mu + \log(1 + \exp(\rho)) \circ \epsilon$ where $\circ$ is point-wise multiplication.

However, in your question you wrote
\begin{align}
\text{prior} &= \log(q(\mathbf{w}|\mu,\rho)) = \sum_i \log(p(w_i | 0, 1)) \\
\text{posterior} &= \log(P(\mathbf{w})) = \sum_i \log(p(w_i | \mu, \sigma^2)) \\
\text{likelihood} &= \log(P(\mathcal{D}|\mathbf{w})) = y \cdot \log(\text{softmax}(\hat{y}))
\end{align}
which does not seem to be correct, since $q(\mathbf{w}|\mu,\rho)$ is the posterior according to the author. Here is what I think the authors meant
\begin{align}
\text{posterior} &= q(\mathbf{w}|\theta) \\
\text{prior} &= P(\mathbf{w}) \\
\text{likelihood} &= P(\mathcal{D}|\mathbf{w})
\end{align}
Also, the authors did not specify the prior $P(\mathbf{w})$ until equation 7 in section 3.3, which is a mixture of Gaussians and not a standard Gaussian as you wrote. Just wanted to point this out.

Here is a rough outline of what you can do in torch to implement this:
1. Create the dataset $\mathcal{D}$ consisting of the vectors $\mathbf{x}_1,\mathbf{x}_2,\dots,\mathbf{x}_N$ and the labels $y_1,y_2,...,y_N$. This can be done by sampling $\mathbf{x}$ vectors from some multivariate distribution, and sampling $y$ labels from a Bernoulli distribution.
2. Sample $\epsilon$ from the multivariate normal distribution $N(0,I)$.
3. Sample initial values for $\mu$ and $\rho$ <s>from the multivariate normal distribution $N(0,I)$ </s>. The authors do not say anything about whether $\mu$ and $\rho$ are themselves random variables with an associated distribution $p(\theta)$ or are treated as unknown constants. For simplicity, I will assume they are treated as unknown constants. In that case, their initial values can be sampled from any distribution. For simplicity, you can sample them using the `torch_randn` function. You only need to do this once to start performing gradient descent. Note that $\mu$ is a column vector and $\rho$ is <s>a diagonal matrix</s> also a column vector because the authors use the element-wise multiplication operation $\circ$ in their paper to multiply $\epsilon$ by $\log(1 + \exp(\rho))$.
4. Compute $\mathbf{w} = \mu + \log(1 + \exp(\rho)) \circ \epsilon$.
5. Compute $\log(q(\mathbf{w}|\mu,\rho))$ by inputting the $\mathbf{w}$,$\mu$, and $\rho$ that you obtained in steps 3 and 4 into the multivariate normal probability density function.
6. Compute $\log(P(\mathbf{w}))$ by inputting the $\mathbf{w}$ that you obtained in step 3 into the multivariate normal probability density function with mean $\mathbf{0}$ and covariance $I$. **Note** again that this does not follow what the authors did in the paper, and that you would instead need to implement equation 7 in the paper.
7. Compute $\log(P(\mathcal{D}|\mathbf{w}))$. Note that you did not specify what $\hat{y}$ is. For the sake of simplicity, I will assume that $\hat{y} = \mathbf{w}^T \mathbf{x}$. This means that $\log(P(\mathcal{D}|\mathbf{w})) = -\sum_{i=1}^N y_i \cdot \log(\hat{y}_i)$. This is just the [cross-entropy](https://en.wikipedia.org/wiki/Cross_entropy) between the $y$ labels and the $\hat{y}$ labels. Note that each $y$ and $\hat{y}$ **must** range between 0 and 1.
8. Compute $f(\mathbf{w},\theta) = \log(q(\mathbf{w}|\theta)) - \log(P(\mathbf{w})) - \log(P(\mathcal{D}|\mathbf{w}))$.
9. Compute the gradients as given in steps 5 and 6 of the algorithm.
10. Update $\mu$ and $\rho$ as given in step 7 of the algorithm.
11. Repeat the following aforementioned steps until convergence: step 2 (sampling $\epsilon$) $\rightarrow$ step 4 (computing $\mathbf{w}$) $\rightarrow$ steps 5,6,7,8,9, and 10.

I am aware that I did not mention how to compute the gradients in torch. If that is still something that you are not sure about, let me know.

## Response to comments

> What is still unclear to me is how to design the hidden layers of neurons. Am i correct in understanding that if i have 100 observations, for (2) this means i sample 100 $\epsilon$'s from a mv normal distribution with $\mu=0$ and a 100x100 identity matrix? for (3) i sample 100 $\mu$'s as a column vector and $\rho$ is 100x100 diagonal matrix of 100 samples of the same mv normal distribution? and for (4) i will get 100 samples $\mathbf{w}$? For my hidden layer i have specified 32 neurons. Where do these come into play?

Suppose you have $N$ observations in your dataset $\mathcal{D}$ of $(\mathbf{x},y)$ pairs, as discussed in step 1 above, and suppose that each $\mathbf{x}$ is $K$-dimensional such that
$$
\mathbf{x} =
\begin{bmatrix}
x_1 \\
x_2 \\
\vdots \\
x_K
\end{bmatrix}
$$
Also, for the sake of simplicity, suppose that your network consists of a single layer and the output of your network is
$$
\hat{y} = \phi(\mathbf{W}^T \mathbf{x} + \mathbf{b})
$$
where $\mathbf{W}$ is a $K \times M$ matrix of weights, where $M$ is the number of possible classes, $\mathbf{b}$ is a $M \times 1$ column vector, and $\phi(\cdot)$ is the softmax function. Recall that the authors wrote
$$
\mathbf{w} = \mu + \log(1 + \exp(\rho)) \circ \epsilon
$$
This be generalized to a $K \times M$ matrix $\mathbf{W}$ by sampling a $K \times M$ matrix $\epsilon$ using `torch_randn`, sampling **initial values** (only need to sample once) of $K \times M$ matrices $\mu$ and $\rho$ using `torch_randn`. Since $\epsilon \sim N(0,I)$, then you can also sample a $K \times M$ matrix $\epsilon$ using `torch_randn`. You can then obtain $\mathbf{W}$ as shown above. Each row of $\mathbf{W}$ represents a single weight vector. You can repeat this process using column vectors for $\mu,\rho,$ and $\epsilon$ to obtain $\mathbf{b}$:
$$
\mathbf{b} = \mu + \log(1 + \exp(\rho)) \circ \epsilon
$$
You can then compute $\hat{y}$ as shown above.

> For step 7, $\hat{y}$ is supposed to represent the predicted values, which are compared to the ground truth $y$. So i think we're talking about the same thing here. You mention both $y$ and $\hat{y}$ have to range between 0 and 1, does that mean i should apply a sigmoid activation function after the output layer? 

I think we have slightly different definitions of $\hat{y}$. I am defining $\hat{y}$ as
$$
\hat{y} = \phi(\mathbf{W}^T \mathbf{x} + \mathbf{b})
$$
___
I'm not experienced in R, so here is my attempt using Python and PyTorch. You can (almost) translate quickly between the two by substituting `.` in Python with `$` in R. I don't get NaN values, but the `f` loss does not decrease consistently. I commented my code as much as possible for clarity. Hope this helps.

```python
import torch
import math

# define pi
pi = torch.tensor(math.pi)

def log_likelihood(x,mu,rho):
    # need to compute negative log-likelihood and NOT the
    # log-likelihood. This is because we are performing gradient descent.
    # In gradient descent, we are trying to minimize a function.
    # Minimizing the negative log-likelihood function is equivalent to
    # maximizing the log-likelihood function.
    
    # need to re-parameterize sigma to keep it positive. See the first
    # paragraph of section 3.2 in the paper and section 3.1 in the paper
    # for details
    sigma = torch.log(1 + torch.exp(rho))
    
    return -torch.log(sigma * torch.sqrt(2 * pi)) - (0.5 * (torch.div(x-mu,sigma) ** 2))

# uncomment one or the other to choose your activation function
act_func = torch.nn.Sigmoid()
# act_func = torch.nn.ReLU()

# the logarithm of the softmax function. See
# https://pytorch.org/docs/stable/generated/torch.nn.LogSoftmax.html
# for details
log_softmax = torch.nn.LogSoftmax(dim = 0)

learning_rate = 1e-5

# STEP 1
# create the dataset D consisting of vectors x_1,x_2,...,x_N and labels
# y_1,y_2,...,y_N

# input dimensionality (number of input features)
d_in = 3
# dimensionality of the hidden layer
d_hidden = 32
# output dimensionality (number of predicted features)
d_out = 1
# number of observations in training set
n = 100
# number of training iterations
num_iter = 100

# create input data. This is a (N x d_in) matrix, where N is the number of
# observations and d_in is the dimensionality of the input
X = torch.randn(n,d_in)

# generate the ground truth with a bernoulli distribution with probability
# p = 0.6
p = 0.6
y = torch.bernoulli(p * torch.ones((n,)))

# In STEP 2, I will sample the required mu and rho matrices and column
# vectors. Since the authors in the paper do not specify p(mu,rho)in the
# paper, then mu and rho are both treated here as unknown constants and
# not as random variables. In this case, I can sample initial values for
# mu and rho from any distributuon. For simplicity, I will use the uniform
# distribution, which can be done using the torch.rand (not torch.randn)
# function.
#
# Here are the weight matrices and bias vectors I am trying to construct
# in STEP 2:
#   
#   weight matrix for the first layer:   
#       W1 = mu_W1 + log(1 + exp(rho_W1)) * epsilon_W1
#   bias vector for the first layer:   
#       b1 = mu_b1 + log(1 + exp(rho_b1)) * epsilon_b1
#   weight matrix for the second layer:   
#       W2 = mu_W2 + log(1 + exp(rho_W2)) * epsilon_W2
#   bias vector for the second layer:   
#       b2 = mu_b2 + log(1 + exp(rho_b2)) * epsilon_b2

# STEP 2a
# sample mu matrix for the first layer with dimensions d_in x d_hidden,
# where d_hidden is the dimensionality of the hidden layer. This mu
# matrix will be used to construct W1, which is the d_in x d_hidden weight
# matrix for the first layer. Note that since I need to compute the
# gradient of f(W,theta) with respect to mu_W1, then I will need to
# set requires_grad = True for mu_W1.

mu_W1 = torch.randn(d_in,d_hidden,requires_grad = True)

# STEP 2b
# sample rho matrix for the first layer with dimensions d_in x d_hidden,
# where d_hidden is the dimensionality of the hidden layer. This rho
# matrix will be used to construct W1, which is the d_in x d_hidden weight
# matrix for the first layer. Note that since I need to compute the
# gradient of f(W,theta) with respect to rho_W1, then I will need to
# set requires_grad = True for rho_W1.

rho_W1 = torch.randn(d_in,d_hidden,requires_grad = True)

# STEP 2c
# sample mu column vector for the first layer with dimensions d_hidden x 1,
# where d_hidden is the dimensionality of the hidden layer. This mu column vector
# will be used to construct b1, which is the d_hidden x 1 bias
# vector for the first layer. Note that since I need to compute the
# gradient of f(W,theta) with respect to mu_b1, then I will need to
# set requires_grad = True for mu_b1.

mu_b1 = torch.randn(d_hidden,1,requires_grad = True)

# STEP 2d
# sample rho column vector for the first layer with dimensions d_hidden x 1,
# where d_hidden is the dimensionality of the hidden layer. This rho column vector
# will be used to construct b1, which is the d_hidden x 1 bias
# vector for the first layer. Note that since I need to compute the
# gradient of f(W,theta) with respect to rho_b1, then I will need to
# set requires_grad = True for rho_b1.

rho_b1 = torch.randn(d_out,1,requires_grad = True)

# sample the rest of the matrices and column vectors

mu_W2 = torch.randn(d_hidden,d_out,requires_grad = True)
rho_W2 = torch.randn(d_hidden,d_out,requires_grad = True)
mu_b2 = torch.randn(d_out,1,requires_grad = True)
rho_b2 = torch.randn(d_out,1,requires_grad = True)

for i in range(num_iter):
    
    # STEP 3a
    # sample epsilon matrix for the first layer with dimensions d_in x d_hidden,
    # where d_hidden is the dimensionality of the hidden layer. This epsilon
    # matrix will be used to construct W1, which is the d_in x d_hidden weight
    # matrix for the first layer. Note that since I do not need to compute the
    # gradient of f(W,theta) with respect to epsilon_W1, then I don't need to
    # set requires_grad = True for epsilon_W1. This will save memory.
    
    epsilon_W1 = torch.randn(d_in,d_hidden)
    
    # STEP 3b
    # sample epsilon column vector for the first layer with dimensions d_hidden x 1,
    # where d_hidden is the dimensionality of the hidden layer. This epsilon
    # matrix will be used to construct b1, which is the d_hidden x 1 bias vector
    # for the first layer. Note that since I do not need to compute the
    # gradient of f(W,theta) with respect to epsilon_b1, then I don't need to
    # set requires_grad = True for epsilon_b1. This will save memory.
    
    epsilon_b1 = torch.randn(d_hidden,1)
    
    # STEP 3c
    # sample epsilon matrix for the second layer with dimensions d_hidden x d_out,
    # where d_out is the dimensionality of the output layer. This epsilon
    # matrix will be used to construct W2, which is the d_hidden x d_out weight
    # matrix for the second layer. Note that since I do not need to compute the
    # gradient of f(W,theta) with respect to epsilon_W2, then I don't need to
    # set requires_grad = True for epsilon_W2. This will save memory.
    
    epsilon_W2 = torch.randn(d_hidden,d_out)
    
    # STEP 3d
    # sample epsilon column vector for the second layer with dimensions d_out x 1,
    # where d_out is the dimensionality of the output layer. This epsilon
    # column vector will be used to construct b2, which is the d_out x 1 bias
    # vector for the second layer. Note that since I do not need to compute the
    # gradient of f(W,theta) with respect to epsilon_b2, then I don't need to
    # set requires_grad = True for epsilon_b2. This will save memory.
    
    epsilon_b2 = torch.randn(d_out,1)
    
    # STEP 4a
    # compute W1 = mu_W1 + log(1 + exp(rho_W1)) * epsilon_W1
    
    W1 = mu_W1 + torch.mul(
                            torch.log(1 + torch.exp(rho_W1)),
                            epsilon_W1
                          )
    
    # STEP 4b
    # compute b1 = mu_b1 + log(1 + exp(rho_b1)) * epsilon_b1
    
    b1 = mu_b1 + torch.mul(
                            torch.log(1 + torch.exp(rho_b1)),
                            epsilon_b1
                          )
    
    # STEP 4c
    # compute W2 = mu_W2 + log(1 + exp(rho_W2)) * epsilon_W2
    
    W2 = mu_W2 + torch.mul(
                            torch.log(1 + torch.exp(rho_W2)),
                            epsilon_W2
                          )
    
    # STEP 4d
    # compute b2 = mu_b2 + log(1 + exp(rho_b2)) * epsilon_b2
    
    b2 = mu_b2 + torch.mul(
                            torch.log(1 + torch.exp(rho_b2)),
                            epsilon_b2
                          )
    
    # STEP 5a
    # compute log(q(W1|mu_W1,rho_W1))
    
    posterior_W1 = log_likelihood(W1,mu_W1,rho_W1).sum()
    
    # STEP 5b
    # compute log(q(b1|mu_b1,rho_b1))
    
    posterior_b1 = log_likelihood(b1,mu_b1,rho_b1).sum()
    
    # STEP 5c
    # compute log(q(W2|mu_W2,rho_W2))
    
    posterior_W2 = log_likelihood(W2,mu_W2,rho_W2).sum()
    
    # STEP 5d
    # compute log(q(b2|mu_b2,rho_b2))
    
    posterior_b2 = log_likelihood(b2,mu_b2,rho_b2).sum()
    
    # STEP 6a
    # compute log(P(W1)). Note that since sigma = log(1 + exp(rho)), and
    # if sigma = 1, then rho = log(exp(1) - 1)
    
    rho = torch.log(torch.exp(torch.tensor(1.)) - 1)
    prior_W1 = log_likelihood(W1,0,rho).sum()
    
    # STEP 6b
    # compute log(P(b1))
    
    prior_b1 = log_likelihood(b1,0,rho).sum()
    
    # STEP 6c
    # compute log(P(W2))
    
    prior_W2 = log_likelihood(W2,0,rho).sum()
    
    # STEP 6d
    # compute log(P(b2))
    
    prior_b2 = log_likelihood(b2,0,rho).sum()
    
    # STEP 7
    # compute log(P(D|W1,W2,b2,b1)), which is the negative of the
    # cross-entropy between all y's and y_hat's in the dataset D. Recall
    # that the cross-entropy between all y's and y_hat's in a dataset of
    # size N is:
    #
    #   - \sum_{i=1}^N y_i \cdot log(\hat{y}_i)
    
    cross_entropy = 0
    
    # iterte over each row of the X matrix to compute the cross-entropy,
    # since each row is a single x vector. Note that this method is
    # inefficient, since batch matrix multiplication would be more
    # efficient. However, this is for illustrative purposes only.
    
    for j in range(X.shape[0]):
        
        # extract the input vector x and convert it into a column vector
        x = X[j,:].unsqueeze(-1)
        
        # output of first layer. Note that W1.T means transpose of W1
        out1 = act_func(torch.matmul(W1.T,x) + b1)
        
        # output of second layer. Note that the logarithm of the softmax
        # function is computed here. The .squeeze() method is used to
        # remove any extra dimensions
        log_y_hat = log_softmax(torch.matmul(W2.T,out1) + b2).squeeze()
        
        # accumulate cross entropy
        cross_entropy = cross_entropy + (y[j] * log_y_hat)
    
    # STEP 8
    # compute f(w,theta). Note that since W1,b1,W2, and b2 are assumed to 
    # be conditionally independent given their corresponding parameters mu
    # and rho, then:
    #
    #   log(q(W1,b1,W2,b2|mu_W1,rho_W1,mu_W2,rho_b2)) = 
    #           log(q(W1|mu_W1,rho_W1)) + log(q(b1|mu_b1,rho_b1)) +
    #           log(q(W2|mu_W2,rho_W2)) + log(q(b2|mu_b2,rho_b2))
    #
    # Note that log(P(D|W1,W2,b2,b1)) is the cross-entropy computed above.
    
    f = (posterior_W1 + posterior_b1 + posterior_W2 + posterior_b2
       - prior_W1 - prior_b1 - prior_W2 - prior_b2
       + cross_entropy)
    
    # STEP 9
    # compute gradients as shown in steps 5 and 6
    
    f.backward()
    # Delta_mu = (W1.grad + b1.grad + W2.grad + b2.grad
    #           + mu_W1.grad + mu_b1.grad + mu_W2.grad + mu_b2.grad)
    # Delta_rho = ((W1.grad * (epsilon_W1 / (1 + torch.exp(-rho_W1)))
    #            + b1.grad * (epsilon_b1 / (1 + torch.exp(-rho_b1)))
    #            + W2.grad * (epsilon_W2 / (1 + torch.exp(-rho_W2)))
    #            + b2.grad * (epsilon_b2 / (1 + torch.exp(-rho_b2))))
    #            + rho_W1.grad + rho_b1.grad + rho_W2.grad + rho_b2.grad)
    
    # STEP 10
    # update mu and rho using gradient descent
    with torch.no_grad():
        mu_W1 -= mu_W1.grad * learning_rate
        mu_b1 -= mu_b1.grad * learning_rate
        mu_W2 -= mu_W2.grad * learning_rate
        mu_b2 -= mu_b2.grad * learning_rate
        mu_W1.grad.zero_()
        mu_b1.grad.zero_()
        mu_W2.grad.zero_()
        mu_b2.grad.zero_()
        rho_W1 -= rho_W1.grad * learning_rate
        rho_b1 -= rho_b1.grad * learning_rate
        rho_W2 -= rho_W2.grad * learning_rate
        rho_b2 -= rho_b2.grad * learning_rate
        rho_W1.grad.zero_()
        rho_b1.grad.zero_()
        rho_W2.grad.zero_()
        rho_b2.grad.zero_()
```