Hi all!

I am new to numpyro and trying to define my model as shown below. I am using NUTS HMC sampling for learning. I have been following this resource predominantly: Example: Bayesian Neural Network — NumPyro documentation.

Could someone help me understand how to set a unit normal prior over my bias terms (`b1`

and `b2`

) ? Is this correct? I want to ensure my weights and biases are updating once HMC starts.

def model(X, Y, D_H, mask):

`D_X, D_Y = X.shape[1], 1 # Sample first layer (put unit normal priors on all weights) w1 = numpyro.sample("w1", dist.Normal(jnp.zeros((D_X, D_H)), jnp.ones((D_X, D_H)))) # dim. D_X D_H b1 = numpyro.sample("b1", dist.Normal(0, 1), sample_shape=(D_H,)) z1 = nonlin(jnp.matmul(X, w1) + b1) # dim. N D_H <= first layer of activations # Sample final layer of weights and neural network output w2 = numpyro.sample("w2", dist.Normal(jnp.zeros((D_H, D_Y)), jnp.ones((D_H, D_Y)))) # dim. D_H D_Y b2 = numpyro.sample("b2", dist.Normal(0, 1), sample_shape=(D_Y,)) z2 = jnp.matmul(z1, w2) + b2 # dim. N D_Y <= output of the neural network # Put prior on observation noise. Note if X~Gamma(a,b) then 1/X~InvGamma(a,b). prec_obs = numpyro.sample("prec_obs", dist.Gamma(0.1, 0.1)) sigma_obs = 1.0 / jnp.sqrt(prec_obs) # Observed data mean prediction pr_y_given_everything = dist.Normal(z2, sigma_obs) numpyro.sample("Y", pr_y_given_everything, obs=Y) # observation obs param is critical. return pr_y_given_everything`

Thank you!

Ashley