Variational inference/Bayesian neural network: fixed data dims in ann…

…_input/ann_output (pymc-devs#506) * fixed data dims * pre-commit Co-authored-by: Oriol (ZBook) <oriol.abril.pla@gmail.com>
iykat · Jan 25, 2023 · 951c0ef · 951c0ef
1 parent 5288e05
commit 951c0ef
Show file tree

Hide file tree

Showing 2 changed files with 225 additions and 75 deletions.
diff --git a/examples/variational_inference/bayesian_neural_network_advi.ipynb b/examples/variational_inference/bayesian_neural_network_advi.ipynb
diff --git a/examples/variational_inference/bayesian_neural_network_advi.myst.md b/examples/variational_inference/bayesian_neural_network_advi.myst.md
@@ -130,8 +130,8 @@ def construct_nn(ann_input, ann_output):
         # "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)
-        ann_output = pm.Data("ann_output", Y_train, mutable=True)
+        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")
 
         # Weights from input to hidden layer
         weights_in_1 = pm.Normal(
@@ -157,6 +157,7 @@ def construct_nn(ann_input, ann_output):
             act_out,
             observed=ann_output,
             total_size=Y_train.shape[0],  # IMPORTANT for minibatches
+            dims="obs_id",
         )
     return neural_network
 
@@ -340,6 +341,7 @@ You might argue that the above network isn't really deep, but note that we could
 
 - This notebook was originally authored as a [blog post](https://twiecki.github.io/blog/2016/06/01/bayesian-deep-learning/) by Thomas Wiecki in 2016
 - Updated by Chris Fonnesbeck for PyMC v4 in 2022
+- Updated by Oriol Abril-Pla and Earl Bellinger in 2023
 
 ## Watermark