Multi-dimensional input Bayesian Neural Network

1

Hello, first time pymc user here and I am trying to solve a non-linear regression problem using a Bayesian Neural Network. My input data has a shape of (14,8) containing 8 correlated input features, and I am predicting a single output, Y.shape = (14,). My goal is to sample both inputs and outputs from the posterior as a means of generating synthetic data, so I need to make sure the input features are correlated. Where I am having some confusion is in the use of the MvNormal prior to implement my input correlation matrix. The code I currently have runs, but I feel as though the priors for the input weights should use MvNormal as well. However, when I try to change the weights_in_1 layer to a MvNormal dist I have shape compatibility issues. Are there any examples out there of Bayesian neural networks that use MvNormal? Should each of the weights and biases priors within the model use MvNormal? Any help is appreciated, thanks!

def build_neural_network(ann_input, ann_output, n_hidden, n_features, cov_matrix):
    # Initialize random weights between each layer
    init_1 = np.random.randn(n_features, n_hidden).astype(floatX)
    init_out = np.random.randn(n_hidden).astype(floatX)

    # Initialize random biases in each layer
    init_b_1 = np.random.randn(n_hidden).astype(floatX)
    init_b_out = np.random.randn(1).astype(floatX)
    with pm.Model() as neural_network:
        # Weights from input to hidden layer
        weights_in_1 = pm.Normal('w_in_1', mu=0, sd=1,shape=(n_features, n_hidden), testval=init_1)
        bias_1 = pm.Normal('b_1', mu=0, sd=1, shape=(n_hidden),testval=init_b_1)
            
        # Weights from hidden layer to output
        weights_2_out = pm.Normal('w_2_out', mu=0, sd=1, shape=(n_hidden,), testval=init_out)
        bias_out = pm.Normal('b_out', mu=0, sd=1, shape=(1),testval=init_b_out)
        
        # Build neural-network using tanh activation function
        act_1 = pm.math.tanh(pm.math.dot(ann_input, weights_in_1) + bias_1)
        act_out = pm.math.dot(act_1, weights_2_out) + bias_out
        
        # Likelihood
        out = pm.Normal('out', mu=act_out, observed=ann_output)
        input = pm.MvNormal("ann_input", mu = np.zeros(n_features), cov = correlation_matrix, 
        observed=ann_input, total_size=X.shape[1])
    return neural_network

ann_input = theano.shared(X)
ann_output = theano.shared(Y)
neural_network = build_neural_network(ann_input, ann_output, 3, X.shape[1], correlation_matrix)

with neural_network:
    s = theano.shared(pm.floatX(1))
    inference = pm.ADVI(cost_part_grad_scale=s)
    tracker = pm.callbacks.Tracker(
    mean=inference.approx.mean.eval,  # callable that returns mean
    std=inference.approx.std.eval  # callable that returns std
    )
    approx = pm.fit(n=10000, method=inference, callbacks=[tracker])

trace = approx.sample(draws=5000)
pm.traceplot(trace)

ann_input.set_value(X)
ann_output.set_value(Y)
with neural_network:
        ppc = pm.sample_posterior_predictive(trace, samples=1000, progressbar=False)

post_out = ppc['out']
post_in = ppc['ann_input']```