I am implementing a Bayesian NN and I want to solve it with SVI. That’s the BNN model

```
D_Y=1
D_H = n_neurons
D_X = 1
w_ = ()
b_ = ()
# sample first layer (we put unit normal priors on all weights)
W_mean = numpyro.param('W1', np.zeros((D_X, D_H)))
W_std = numpyro.param('B1', np.ones((D_X, D_H)))
B_mean = numpyro.param('W1', np.zeros((1, D_H)))
B_std = numpyro.param('B1', np.ones((1, D_H)))
w = numpyro.sample(f"w{1}", dist.Normal(W_mean, W_std))
b = numpyro.sample(f"b{1}", dist.Normal(B_mean, B_std))
w_ = w_ + (w,)
b_ = b_ + (b,)
for i in range(1,n_layers-1):
W_mean = numpyro.param(f'W{i+1}', np.zeros((D_H, D_H)))
W_std = numpyro.param(f'B{i+1}', np.ones((D_H, D_H)))
pdb.set_trace()
B_mean = numpyro.param(f'W{i+1}', np.zeros((1, D_H)))
B_std = numpyro.param(f'B{i+1}', np.ones((1, D_H)))
w = numpyro.sample(f"w{i+1}", dist.Normal(W_mean, W_std))
b = numpyro.sample(f"b{i+1}", dist.Normal(B_mean, B_std))
w_ = w_ + (w,)
b_ = b_ + (b,)
W_mean = numpyro.param(f'W{n_layers}', np.zeros((D_H, D_Y)))
W_std = numpyro.param(f'B{n_layers}', np.ones((D_H, D_Y)))
B_mean = numpyro.param(f'W{n_layers}', np.zeros((1, D_Y)))
B_std = numpyro.param(f'B{n_layers}', np.ones((1, D_Y)))
w = numpyro.sample(f"w{n_layers}", dist.Normal(W_mean, W_std))
b = numpyro.sample(f"b{n_layers}", dist.Normal(B_mean, B_std))
w_ = w_ + (w,)
b_ = b_ + (b,)
theta = (w_,b_)
```

The error is caused by the incorrect size of the weights matrix “w2” inside the loop.If I put a breakpoint just after `W_std`

I see that the first time the model is called w2 has the correct dimension 2x2, but the second time, just before the inference is solved is 2x1, although it is clear that D_H = 2.

I do not understand how this is possible, since I clearly stated it should be of dimension 2x2. What am I missing?

.