Variational inference/Bayesian neural network: fixed data dims in ann_input/ann_output

Fixes issue https://github.com/pymc-devs/pymc-examples/issues/504 using the solution given by OriolAbril, i.e., making the following change:

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")
...
# Binary classification -> Bernoulli likelihood
out = pm.Bernoulli(
    "out",
    act_out,
    observed=ann_output,
    total_size=Y_train.shape[0],  # IMPORTANT for minibatches
    dims="obs_id",
)

The notebook now runs to completion.

pymc-devs / pymc-examples

Variational inference/Bayesian neural network: fixed data dims in ann_input/ann_output #506