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

[earlbellinger]

Fixes issue #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.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[OriolAbril]

Thanks!