Strange error during SVI

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?
.

I just realized that I completely messed up the variable names :sweat_smile:. Here is the working 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_mean', np.zeros((D_X, D_H)))
        W_std = numpyro.param('W1_std', np.ones((D_X, D_H)))
        B_mean = numpyro.param('B1_mean', np.zeros((1, D_H)))
        B_std = numpyro.param('B1_std', 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}_mean', np.zeros((D_H, D_H)))
            W_std = numpyro.param(f'W{i+1}_std', np.ones((D_H, D_H)))
            B_mean = numpyro.param(f'B{i+1}_mean', np.zeros((1, D_H)))
            B_std = numpyro.param(f'B{i+1}_std', 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}_mean', np.zeros((D_H, D_Y)))
        W_std = numpyro.param(f'W{n_layers}_std', np.ones((D_H, D_Y)))
        B_mean = numpyro.param(f'B{n_layers}_mean', np.zeros((1, D_Y)))
        B_std = numpyro.param(f'B{n_layers}_std', 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_)