Bug in "Variational Inference: Bayesian Neural Networks" notebook

[yoshida-chem]

Notebook title: Variational Inference: Bayesian Neural Networks
Notebook url: https://www.pymc.io/projects/examples/en/latest/variational_inference/bayesian_neural_network_advi.html

Issue description

One needs to pass all the variables to a single Minibatch call so that their slices are identical, otherwise it would randomly pair some X with unrelated y.

minibatch_x = pm.Minibatch(X_train, batch_size=50)
minibatch_y = pm.Minibatch(Y_train, batch_size=50)

ann_input and ann_output are overwritten by X _train and Y_train.

def construct_nn(ann_input, ann_output):
    n_hidden = 5

    # Initialize random weights between each layer
    init_1 = rng.standard_normal(size=(X_train.shape[1], n_hidden)).astype(floatX)
    init_2 = rng.standard_normal(size=(n_hidden, n_hidden)).astype(floatX)
    init_out = rng.standard_normal(size=n_hidden).astype(floatX)

    coords = {
        "hidden_layer_1": np.arange(n_hidden),
        "hidden_layer_2": np.arange(n_hidden),
        "train_cols": np.arange(X_train.shape[1]),
        # "obs_id": np.arange(X_train.shape[0]),
    }
    with pm.Model(coords=coords) as neural_network:
        ann_input = pm.Data("ann_input", X_train, mutable=True, dims=("obs_id", "train_cols"))
        ann_output = pm.Data("ann_output", Y_train, mutable=True, dims="obs_id")

Proposed solution

One needs to pass all the variables to a single Minibatch call so that their slices are identical.

minibatch_x, minibatch_y = pm.Minibatch(X_train, Y_train, batch_size=50)

[Deepakchowdavarapu]

@yoshida-chem , @twiecki @ can you review my PR #689

[Deepakchowdavarapu]

@yoshida-chem @twiecki @