diff --git a/examples/mixture_models/dependent_density_regression.ipynb b/examples/mixture_models/dependent_density_regression.ipynb
index 0cd667489..c5628c0d5 100644
--- a/examples/mixture_models/dependent_density_regression.ipynb
+++ b/examples/mixture_models/dependent_density_regression.ipynb
@@ -2,16 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "metadata": {
-    "papermill": {
-     "duration": 0.021485,
-     "end_time": "2020-12-18T15:39:09.385692",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:09.364207",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "# Dependent density regression\n",
     "In another [example](dp_mix.ipynb), we showed how to use Dirichlet processes to perform Bayesian nonparametric density estimation.  This example expands on the previous one, illustrating dependent density regression.\n",
@@ -21,29 +12,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:39:09.433230Z",
-     "iopub.status.busy": "2020-12-18T15:39:09.432485Z",
-     "iopub.status.idle": "2020-12-18T15:39:16.049724Z",
-     "shell.execute_reply": "2020-12-18T15:39:16.048768Z"
-    },
-    "papermill": {
-     "duration": 6.644193,
-     "end_time": "2020-12-18T15:39:16.049848",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:09.405655",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 2,
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Running on PyMC3 v3.10.0\n"
+      "Running on PyMC3 v3.11.2\n"
      ]
     }
    ],
@@ -64,23 +40,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:39:16.118903Z",
-     "iopub.status.busy": "2020-12-18T15:39:16.116655Z",
-     "iopub.status.idle": "2020-12-18T15:39:16.119741Z",
-     "shell.execute_reply": "2020-12-18T15:39:16.120266Z"
-    },
-    "papermill": {
-     "duration": 0.048777,
-     "end_time": "2020-12-18T15:39:16.120430",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:16.071653",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 3,
+   "metadata": {},
    "outputs": [],
    "source": [
     "%config InlineBackend.figure_format = 'retina'\n",
@@ -93,39 +54,15 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "papermill": {
-     "duration": 0.020257,
-     "end_time": "2020-12-18T15:39:16.160652",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:16.140395",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "We will use the LIDAR data set from Larry Wasserman's excellent book, [_All of Nonparametric Statistics_](http://www.stat.cmu.edu/~larry/all-of-nonpar/).  We standardize the data set to improve the rate of convergence of our samples."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:39:16.210611Z",
-     "iopub.status.busy": "2020-12-18T15:39:16.209819Z",
-     "iopub.status.idle": "2020-12-18T15:39:16.708850Z",
-     "shell.execute_reply": "2020-12-18T15:39:16.708223Z"
-    },
-    "papermill": {
-     "duration": 0.527789,
-     "end_time": "2020-12-18T15:39:16.708982",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:16.181193",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 4,
+   "metadata": {},
    "outputs": [],
    "source": [
     "DATA_URI = \"http://www.stat.cmu.edu/~larry/all-of-nonpar/=data/lidar.dat\"\n",
@@ -142,23 +79,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:39:16.761187Z",
-     "iopub.status.busy": "2020-12-18T15:39:16.760225Z",
-     "iopub.status.idle": "2020-12-18T15:39:16.769609Z",
-     "shell.execute_reply": "2020-12-18T15:39:16.770070Z"
-    },
-    "papermill": {
-     "duration": 0.040034,
-     "end_time": "2020-12-18T15:39:16.770185",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:16.730151",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 5,
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -236,7 +158,7 @@
        "4    396 -0.059913  -1.655168      0.818631"
       ]
      },
-     "execution_count": 4,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -247,39 +169,15 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "papermill": {
-     "duration": 0.021121,
-     "end_time": "2020-12-18T15:39:16.811381",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:16.790260",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "We plot the LIDAR data below."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:39:16.858458Z",
-     "iopub.status.busy": "2020-12-18T15:39:16.857709Z",
-     "iopub.status.idle": "2020-12-18T15:39:17.276172Z",
-     "shell.execute_reply": "2020-12-18T15:39:17.276774Z"
-    },
-    "papermill": {
-     "duration": 0.444885,
-     "end_time": "2020-12-18T15:39:17.276934",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:16.832049",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 6,
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -311,16 +209,7 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "papermill": {
-     "duration": 0.022903,
-     "end_time": "2020-12-18T15:39:17.324348",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:17.301445",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "This data set has a two interesting properties that make it useful for illustrating dependent density regression.\n",
     "\n",
@@ -332,23 +221,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:39:17.394681Z",
-     "iopub.status.busy": "2020-12-18T15:39:17.393643Z",
-     "iopub.status.idle": "2020-12-18T15:39:18.315621Z",
-     "shell.execute_reply": "2020-12-18T15:39:18.314834Z"
-    },
-    "papermill": {
-     "duration": 0.968079,
-     "end_time": "2020-12-18T15:39:18.315769",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:17.347690",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 7,
+   "metadata": {},
    "outputs": [],
    "source": [
     "fig, (scatter_ax, hist_ax) = plt.subplots(ncols=2, figsize=(16, 6))\n",
@@ -392,23 +266,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:39:18.367248Z",
-     "iopub.status.busy": "2020-12-18T15:39:18.366431Z",
-     "iopub.status.idle": "2020-12-18T15:39:22.854668Z",
-     "shell.execute_reply": "2020-12-18T15:39:22.853747Z"
-    },
-    "papermill": {
-     "duration": 4.515607,
-     "end_time": "2020-12-18T15:39:22.854784",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:18.339177",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 8,
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -1442,7 +1301,7 @@
        "<IPython.core.display.HTML object>"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1453,16 +1312,7 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "papermill": {
-     "duration": 0.026205,
-     "end_time": "2020-12-18T15:39:22.908143",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:22.881938",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "As we slice the data with a window sliding along the x-axis in the left plot, the empirical distribution of the y-values of the points in the window varies in the right plot.  An important aspect of this approach is that the density estimates that correspond to close values of the predictor are similar.\n",
     "\n",
@@ -1489,23 +1339,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:39:22.972869Z",
-     "iopub.status.busy": "2020-12-18T15:39:22.970896Z",
-     "iopub.status.idle": "2020-12-18T15:39:22.973674Z",
-     "shell.execute_reply": "2020-12-18T15:39:22.974219Z"
-    },
-    "papermill": {
-     "duration": 0.038654,
-     "end_time": "2020-12-18T15:39:22.974358",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:22.935704",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 9,
+   "metadata": {},
    "outputs": [],
    "source": [
     "def norm_cdf(z):\n",
@@ -1520,51 +1355,27 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:39:23.038435Z",
-     "iopub.status.busy": "2020-12-18T15:39:23.037472Z",
-     "iopub.status.idle": "2020-12-18T15:39:53.434884Z",
-     "shell.execute_reply": "2020-12-18T15:39:53.433651Z"
-    },
-    "papermill": {
-     "duration": 30.433368,
-     "end_time": "2020-12-18T15:39:53.435014",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:23.001646",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 10,
+   "metadata": {},
    "outputs": [],
    "source": [
-    "N, _ = df.shape\n",
+    "N = len(df)\n",
     "K = 20\n",
     "\n",
     "std_range = df.std_range.values[:, np.newaxis]\n",
     "std_logratio = df.std_logratio.values\n",
     "\n",
-    "with pm.Model() as model:\n",
-    "    alpha = pm.Normal(\"alpha\", 0.0, 5.0, shape=K)\n",
-    "    beta = pm.Normal(\"beta\", 0.0, 5.0, shape=(1, K))\n",
+    "with pm.Model(coords={\"N\": np.arange(N), \"K\": np.arange(K) + 1, \"one\": [1]}) as model:\n",
+    "    alpha = pm.Normal(\"alpha\", 0.0, 5.0, dims=\"K\")\n",
+    "    beta = pm.Normal(\"beta\", 0.0, 5.0, dims=(\"one\", \"K\"))\n",
     "    x = pm.Data(\"x\", std_range)\n",
     "    v = norm_cdf(alpha + pm.math.dot(x, beta))\n",
-    "    w = pm.Deterministic(\"w\", stick_breaking(v))"
+    "    w = pm.Deterministic(\"w\", stick_breaking(v), dims=[\"N\", \"K\"])"
    ]
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "papermill": {
-     "duration": 0.026426,
-     "end_time": "2020-12-18T15:39:53.491766",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:53.465340",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "We have defined `x` as a `pm.Data` container in order to use `PyMC3`'s posterior prediction capabilities later.\n",
     "\n",
@@ -1588,66 +1399,27 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:39:53.557290Z",
-     "iopub.status.busy": "2020-12-18T15:39:53.556204Z",
-     "iopub.status.idle": "2020-12-18T15:39:53.787741Z",
-     "shell.execute_reply": "2020-12-18T15:39:53.787110Z"
-    },
-    "papermill": {
-     "duration": 0.26821,
-     "end_time": "2020-12-18T15:39:53.787869",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:53.519659",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 11,
+   "metadata": {},
    "outputs": [],
    "source": [
     "with model:\n",
-    "    gamma = pm.Normal(\"gamma\", 0.0, 10.0, shape=K)\n",
-    "    delta = pm.Normal(\"delta\", 0.0, 10.0, shape=(1, K))\n",
+    "    gamma = pm.Normal(\"gamma\", 0.0, 10.0, dims=\"K\")\n",
+    "    delta = pm.Normal(\"delta\", 0.0, 10.0, dims=(\"one\", \"K\"))\n",
     "    mu = pm.Deterministic(\"mu\", gamma + pm.math.dot(x, delta))"
    ]
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "papermill": {
-     "duration": 0.030827,
-     "end_time": "2020-12-18T15:39:53.849720",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:53.818893",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "Finally, we place the prior $\\tau_i \\sim \\textrm{Gamma}(1, 1)$ on the component precisions."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:39:53.921597Z",
-     "iopub.status.busy": "2020-12-18T15:39:53.920393Z",
-     "iopub.status.idle": "2020-12-18T15:40:49.643976Z",
-     "shell.execute_reply": "2020-12-18T15:40:49.644515Z"
-    },
-    "papermill": {
-     "duration": 55.765002,
-     "end_time": "2020-12-18T15:40:49.644639",
-     "exception": false,
-     "start_time": "2020-12-18T15:39:53.879637",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 12,
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -1697,7 +1469,7 @@
        "<text text-anchor=\"middle\" x=\"59\" y=\"-368.73\" font-family=\"Times,serif\" font-size=\"14.00\">Gamma</text>\n",
        "</g>\n",
        "<!-- obs -->\n",
-       "<g id=\"node10\" class=\"node\">\n",
+       "<g id=\"node9\" class=\"node\">\n",
        "<title>obs</title>\n",
        "<ellipse fill=\"lightgrey\" stroke=\"black\" cx=\"363\" cy=\"-165.48\" rx=\"72.66\" ry=\"37.45\"/>\n",
        "<text text-anchor=\"middle\" x=\"363\" y=\"-176.78\" font-family=\"Times,serif\" font-size=\"14.00\">obs</text>\n",
@@ -1705,35 +1477,13 @@
        "<text text-anchor=\"middle\" x=\"363\" y=\"-146.78\" font-family=\"Times,serif\" font-size=\"14.00\">NormalMixture</text>\n",
        "</g>\n",
        "<!-- tau&#45;&gt;obs -->\n",
-       "<g id=\"edge8\" class=\"edge\">\n",
+       "<g id=\"edge10\" class=\"edge\">\n",
        "<title>tau&#45;&gt;obs</title>\n",
        "<path fill=\"none\" stroke=\"black\" d=\"M79.85,-354.5C88.59,-342.51 99.42,-329.3 111,-318.95 169.04,-267.1 247.13,-223.01 300.93,-195.78\"/>\n",
        "<polygon fill=\"black\" stroke=\"black\" points=\"302.64,-198.84 310.01,-191.22 299.5,-192.58 302.64,-198.84\"/>\n",
        "</g>\n",
-       "<!-- gamma -->\n",
-       "<g id=\"node2\" class=\"node\">\n",
-       "<title>gamma</title>\n",
-       "<ellipse fill=\"none\" stroke=\"black\" cx=\"162\" cy=\"-387.43\" rx=\"41.94\" ry=\"37.45\"/>\n",
-       "<text text-anchor=\"middle\" x=\"162\" y=\"-398.73\" font-family=\"Times,serif\" font-size=\"14.00\">gamma</text>\n",
-       "<text text-anchor=\"middle\" x=\"162\" y=\"-383.73\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
-       "<text text-anchor=\"middle\" x=\"162\" y=\"-368.73\" font-family=\"Times,serif\" font-size=\"14.00\">Normal</text>\n",
-       "</g>\n",
-       "<!-- mu -->\n",
-       "<g id=\"node8\" class=\"node\">\n",
-       "<title>mu</title>\n",
-       "<polygon fill=\"none\" stroke=\"black\" points=\"408.5,-302.95 317.5,-302.95 317.5,-249.95 408.5,-249.95 408.5,-302.95\"/>\n",
-       "<text text-anchor=\"middle\" x=\"363\" y=\"-287.75\" font-family=\"Times,serif\" font-size=\"14.00\">mu</text>\n",
-       "<text text-anchor=\"middle\" x=\"363\" y=\"-272.75\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
-       "<text text-anchor=\"middle\" x=\"363\" y=\"-257.75\" font-family=\"Times,serif\" font-size=\"14.00\">Deterministic</text>\n",
-       "</g>\n",
-       "<!-- gamma&#45;&gt;mu -->\n",
-       "<g id=\"edge6\" class=\"edge\">\n",
-       "<title>gamma&#45;&gt;mu</title>\n",
-       "<path fill=\"none\" stroke=\"black\" d=\"M180.02,-353.42C188.28,-341.01 199.13,-327.77 212,-318.95 240.21,-299.63 277.15,-289.27 307.44,-283.74\"/>\n",
-       "<polygon fill=\"black\" stroke=\"black\" points=\"308.15,-287.17 317.43,-282.04 306.98,-280.27 308.15,-287.17\"/>\n",
-       "</g>\n",
        "<!-- alpha -->\n",
-       "<g id=\"node3\" class=\"node\">\n",
+       "<g id=\"node2\" class=\"node\">\n",
        "<title>alpha</title>\n",
        "<ellipse fill=\"none\" stroke=\"black\" cx=\"263\" cy=\"-387.43\" rx=\"41.94\" ry=\"37.45\"/>\n",
        "<text text-anchor=\"middle\" x=\"263\" y=\"-398.73\" font-family=\"Times,serif\" font-size=\"14.00\">alpha</text>\n",
@@ -1749,11 +1499,33 @@
        "<text text-anchor=\"middle\" x=\"472\" y=\"-257.75\" font-family=\"Times,serif\" font-size=\"14.00\">Deterministic</text>\n",
        "</g>\n",
        "<!-- alpha&#45;&gt;w -->\n",
-       "<g id=\"edge3\" class=\"edge\">\n",
+       "<g id=\"edge1\" class=\"edge\">\n",
        "<title>alpha&#45;&gt;w</title>\n",
        "<path fill=\"none\" stroke=\"black\" d=\"M281.19,-353.52C290.18,-340.48 302.29,-326.75 317,-318.95 336.89,-308.41 396.61,-317.97 418,-310.95 420.94,-309.99 423.89,-308.83 426.81,-307.53\"/>\n",
        "<polygon fill=\"black\" stroke=\"black\" points=\"428.57,-310.57 435.97,-302.99 425.46,-304.3 428.57,-310.57\"/>\n",
        "</g>\n",
+       "<!-- gamma -->\n",
+       "<g id=\"node3\" class=\"node\">\n",
+       "<title>gamma</title>\n",
+       "<ellipse fill=\"none\" stroke=\"black\" cx=\"162\" cy=\"-387.43\" rx=\"41.94\" ry=\"37.45\"/>\n",
+       "<text text-anchor=\"middle\" x=\"162\" y=\"-398.73\" font-family=\"Times,serif\" font-size=\"14.00\">gamma</text>\n",
+       "<text text-anchor=\"middle\" x=\"162\" y=\"-383.73\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
+       "<text text-anchor=\"middle\" x=\"162\" y=\"-368.73\" font-family=\"Times,serif\" font-size=\"14.00\">Normal</text>\n",
+       "</g>\n",
+       "<!-- mu -->\n",
+       "<g id=\"node8\" class=\"node\">\n",
+       "<title>mu</title>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"408.5,-302.95 317.5,-302.95 317.5,-249.95 408.5,-249.95 408.5,-302.95\"/>\n",
+       "<text text-anchor=\"middle\" x=\"363\" y=\"-287.75\" font-family=\"Times,serif\" font-size=\"14.00\">mu</text>\n",
+       "<text text-anchor=\"middle\" x=\"363\" y=\"-272.75\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
+       "<text text-anchor=\"middle\" x=\"363\" y=\"-257.75\" font-family=\"Times,serif\" font-size=\"14.00\">Deterministic</text>\n",
+       "</g>\n",
+       "<!-- gamma&#45;&gt;mu -->\n",
+       "<g id=\"edge5\" class=\"edge\">\n",
+       "<title>gamma&#45;&gt;mu</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M180.02,-353.42C188.28,-341.01 199.13,-327.77 212,-318.95 240.21,-299.63 277.15,-289.27 307.44,-283.74\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"308.15,-287.17 317.43,-282.04 306.98,-280.27 308.15,-287.17\"/>\n",
+       "</g>\n",
        "<!-- beta -->\n",
        "<g id=\"node4\" class=\"node\">\n",
        "<title>beta</title>\n",
@@ -1763,7 +1535,7 @@
        "<text text-anchor=\"middle\" x=\"472\" y=\"-368.73\" font-family=\"Times,serif\" font-size=\"14.00\">Normal</text>\n",
        "</g>\n",
        "<!-- beta&#45;&gt;w -->\n",
-       "<g id=\"edge2\" class=\"edge\">\n",
+       "<g id=\"edge3\" class=\"edge\">\n",
        "<title>beta&#45;&gt;w</title>\n",
        "<path fill=\"none\" stroke=\"black\" d=\"M472,-349.84C472,-338.1 472,-325.11 472,-313.42\"/>\n",
        "<polygon fill=\"black\" stroke=\"black\" points=\"475.5,-313.11 472,-303.11 468.5,-313.11 475.5,-313.11\"/>\n",
@@ -1777,7 +1549,7 @@
        "<text text-anchor=\"middle\" x=\"371\" y=\"-368.73\" font-family=\"Times,serif\" font-size=\"14.00\">Normal</text>\n",
        "</g>\n",
        "<!-- delta&#45;&gt;mu -->\n",
-       "<g id=\"edge5\" class=\"edge\">\n",
+       "<g id=\"edge6\" class=\"edge\">\n",
        "<title>delta&#45;&gt;mu</title>\n",
        "<path fill=\"none\" stroke=\"black\" d=\"M368.31,-349.84C367.45,-338.1 366.5,-325.11 365.64,-313.42\"/>\n",
        "<polygon fill=\"black\" stroke=\"black\" points=\"369.11,-312.83 364.88,-303.11 362.13,-313.34 369.11,-312.83\"/>\n",
@@ -1791,7 +1563,7 @@
        "<text text-anchor=\"middle\" x=\"565\" y=\"-368.73\" font-family=\"Times,serif\" font-size=\"14.00\">Data</text>\n",
        "</g>\n",
        "<!-- x&#45;&gt;w -->\n",
-       "<g id=\"edge1\" class=\"edge\">\n",
+       "<g id=\"edge2\" class=\"edge\">\n",
        "<title>x&#45;&gt;w</title>\n",
        "<path fill=\"none\" stroke=\"black\" d=\"M553.24,-360.85C546.4,-347.53 536.97,-331.47 526,-318.95 523.24,-315.8 520.21,-312.73 517.03,-309.77\"/>\n",
        "<polygon fill=\"black\" stroke=\"black\" points=\"519.25,-307.06 509.43,-303.1 514.63,-312.33 519.25,-307.06\"/>\n",
@@ -1803,19 +1575,19 @@
        "<polygon fill=\"black\" stroke=\"black\" points=\"410.25,-304.26 399.73,-303.02 407.18,-310.55 410.25,-304.26\"/>\n",
        "</g>\n",
        "<!-- w&#45;&gt;obs -->\n",
-       "<g id=\"edge9\" class=\"edge\">\n",
+       "<g id=\"edge8\" class=\"edge\">\n",
        "<title>w&#45;&gt;obs</title>\n",
        "<path fill=\"none\" stroke=\"black\" d=\"M448.05,-249.9C438.77,-240.12 427.98,-228.93 418,-218.95 413.75,-214.71 409.28,-210.32 404.79,-205.96\"/>\n",
        "<polygon fill=\"black\" stroke=\"black\" points=\"406.9,-203.13 397.27,-198.71 402.04,-208.17 406.9,-203.13\"/>\n",
        "</g>\n",
        "<!-- mu&#45;&gt;obs -->\n",
-       "<g id=\"edge10\" class=\"edge\">\n",
+       "<g id=\"edge9\" class=\"edge\">\n",
        "<title>mu&#45;&gt;obs</title>\n",
        "<path fill=\"none\" stroke=\"black\" d=\"M363,-249.9C363,-238.99 363,-225.9 363,-213.36\"/>\n",
        "<polygon fill=\"black\" stroke=\"black\" points=\"366.5,-213.01 363,-203.01 359.5,-213.01 366.5,-213.01\"/>\n",
        "</g>\n",
        "<!-- y -->\n",
-       "<g id=\"node9\" class=\"node\">\n",
+       "<g id=\"node10\" class=\"node\">\n",
        "<title>y</title>\n",
        "<path fill=\"lightgrey\" stroke=\"black\" d=\"M378,-92C378,-92 348,-92 348,-92 342,-92 336,-86 336,-80 336,-80 336,-51 336,-51 336,-45 342,-39 348,-39 348,-39 378,-39 378,-39 384,-39 390,-45 390,-51 390,-51 390,-80 390,-80 390,-86 384,-92 378,-92\"/>\n",
        "<text text-anchor=\"middle\" x=\"363\" y=\"-76.8\" font-family=\"Times,serif\" font-size=\"14.00\">y</text>\n",
@@ -1832,17 +1604,17 @@
        "</svg>\n"
       ],
       "text/plain": [
-       "<graphviz.dot.Digraph at 0x7f4ac8714390>"
+       "<graphviz.dot.Digraph at 0x7f34cf88ed90>"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "with model:\n",
-    "    tau = pm.Gamma(\"tau\", 1.0, 1.0, shape=K)\n",
+    "    tau = pm.Gamma(\"tau\", 1.0, 1.0, dims=\"K\")\n",
     "    y = pm.Data(\"y\", std_logratio)\n",
     "    obs = pm.NormalMixture(\"obs\", w, mu, tau=tau, observed=y)\n",
     "\n",
@@ -1851,39 +1623,15 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "papermill": {
-     "duration": 0.027851,
-     "end_time": "2020-12-18T15:40:49.700524",
-     "exception": false,
-     "start_time": "2020-12-18T15:40:49.672673",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "We now sample from the dependent density regression model."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:40:49.763101Z",
-     "iopub.status.busy": "2020-12-18T15:40:49.762286Z",
-     "iopub.status.idle": "2020-12-18T15:48:50.089080Z",
-     "shell.execute_reply": "2020-12-18T15:48:50.087914Z"
-    },
-    "papermill": {
-     "duration": 480.361556,
-     "end_time": "2020-12-18T15:48:50.089232",
-     "exception": false,
-     "start_time": "2020-12-18T15:40:49.727676",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 13,
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -1903,7 +1651,7 @@
        "            }\n",
        "        </style>\n",
        "      <progress value='30000' class='' max='30000' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
-       "      100.00% [30000/30000 03:12<00:00 Sampling chain 0, 0 divergences]\n",
+       "      100.00% [30000/30000 02:56<00:00 Sampling chain 0, 0 divergences]\n",
        "    </div>\n",
        "    "
       ],
@@ -1932,7 +1680,7 @@
        "            }\n",
        "        </style>\n",
        "      <progress value='30000' class='' max='30000' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
-       "      100.00% [30000/30000 03:06<00:00 Sampling chain 1, 0 divergences]\n",
+       "      100.00% [30000/30000 02:50<00:00 Sampling chain 1, 0 divergences]\n",
        "    </div>\n",
        "    "
       ],
@@ -1950,79 +1698,20 @@
     "\n",
     "with model:\n",
     "    step = pm.Metropolis()\n",
-    "    trace = pm.sample(SAMPLES, step=step, tune=BURN, random_seed=SEED, return_inferencedata=True)"
+    "    trace = pm.sample(SAMPLES, tune=BURN, step=step, random_seed=SEED, return_inferencedata=True)"
    ]
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "papermill": {
-     "duration": 0.029328,
-     "end_time": "2020-12-18T15:48:50.149503",
-     "exception": false,
-     "start_time": "2020-12-18T15:48:50.120175",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "To verify that truncation did not unduly influence our results, we plot the largest posterior expected mixture weight for each component.  (In this model, each point has a mixture weight for each component, so we plot the maximum mixture weight for each component across all data points in order to judge if the component exerts any influence on the posterior.)"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:48:50.215378Z",
-     "iopub.status.busy": "2020-12-18T15:48:50.214431Z",
-     "iopub.status.idle": "2020-12-18T15:48:50.218196Z",
-     "shell.execute_reply": "2020-12-18T15:48:50.218804Z"
-    },
-    "papermill": {
-     "duration": 0.040946,
-     "end_time": "2020-12-18T15:48:50.218924",
-     "exception": false,
-     "start_time": "2020-12-18T15:48:50.177978",
-     "status": "completed"
-    },
-    "tags": []
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "(2, 20000, 221, 20)"
-      ]
-     },
-     "execution_count": 13,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "trace.posterior[\"w\"].shape  # (n_chains, n_samples, N, K)"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 14,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:48:50.285924Z",
-     "iopub.status.busy": "2020-12-18T15:48:50.284921Z",
-     "iopub.status.idle": "2020-12-18T15:48:50.844573Z",
-     "shell.execute_reply": "2020-12-18T15:48:50.845231Z"
-    },
-    "papermill": {
-     "duration": 0.5966,
-     "end_time": "2020-12-18T15:48:50.845425",
-     "exception": false,
-     "start_time": "2020-12-18T15:48:50.248825",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -2043,10 +1732,8 @@
    "source": [
     "fig, ax = plt.subplots(figsize=(8, 6))\n",
     "\n",
-    "w_posterior = trace.posterior[\"w\"].data\n",
-    "# reshape to (n_chains*n_samples, N, K)\n",
-    "w_posterior = w_posterior.reshape(-1, w_posterior.shape[2], w_posterior.shape[3])\n",
-    "ax.bar(np.arange(K) + 1, w_posterior.mean(axis=0).max(axis=0))\n",
+    "max_mixture_weights = trace.posterior[\"w\"].mean((\"chain\", \"draw\")).max(\"N\")\n",
+    "ax.bar(max_mixture_weights.coords.to_index(), max_mixture_weights)\n",
     "\n",
     "ax.set_xlim(1 - 0.5, K + 0.5)\n",
     "ax.set_xticks(np.arange(0, K, 2) + 1)\n",
@@ -2057,16 +1744,7 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "papermill": {
-     "duration": 0.031097,
-     "end_time": "2020-12-18T15:48:50.911683",
-     "exception": false,
-     "start_time": "2020-12-18T15:48:50.880586",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "Since only three mixture components have appreciable posterior expected weight for any data point, we can be fairly certain that truncation did not unduly influence our results.  (If most components had appreciable posterior expected weight, truncation may have influenced the results, and we would have increased the number of components and sampled again.)\n",
     "\n",
@@ -2076,28 +1754,13 @@
   {
    "cell_type": "code",
    "execution_count": 15,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:48:51.606803Z",
-     "iopub.status.busy": "2020-12-18T15:48:51.130587Z",
-     "iopub.status.idle": "2020-12-18T15:52:56.720126Z",
-     "shell.execute_reply": "2020-12-18T15:52:56.720612Z"
-    },
-    "papermill": {
-     "duration": 245.777219,
-     "end_time": "2020-12-18T15:52:56.720777",
-     "exception": false,
-     "start_time": "2020-12-18T15:48:50.943558",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/opt/conda/lib/python3.7/site-packages/pymc3/sampling.py:1691: UserWarning: samples parameter is smaller than nchains times ndraws, some draws and/or chains may not be represented in the returned posterior predictive sample\n",
+      "/opt/conda/lib/python3.7/site-packages/pymc3/sampling.py:1690: UserWarning: samples parameter is smaller than nchains times ndraws, some draws and/or chains may not be represented in the returned posterior predictive sample\n",
       "  \"samples parameter is smaller than nchains times ndraws, some draws \"\n"
      ]
     },
@@ -2119,7 +1782,7 @@
        "            }\n",
        "        </style>\n",
        "      <progress value='5000' class='' max='5000' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
-       "      100.00% [5000/5000 03:21<00:00]\n",
+       "      100.00% [5000/5000 03:23<00:00]\n",
        "    </div>\n",
        "    "
       ],
@@ -2143,16 +1806,7 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "papermill": {
-     "duration": 0.032508,
-     "end_time": "2020-12-18T15:52:56.786756",
-     "exception": false,
-     "start_time": "2020-12-18T15:52:56.754248",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "Below we plot the posterior expected value and the 95% posterior credible interval."
    ]
@@ -2160,22 +1814,7 @@
   {
    "cell_type": "code",
    "execution_count": 16,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:52:56.864375Z",
-     "iopub.status.busy": "2020-12-18T15:52:56.863260Z",
-     "iopub.status.idle": "2020-12-18T15:52:57.294728Z",
-     "shell.execute_reply": "2020-12-18T15:52:57.295278Z"
-    },
-    "papermill": {
-     "duration": 0.475588,
-     "end_time": "2020-12-18T15:52:57.295445",
-     "exception": false,
-     "start_time": "2020-12-18T15:52:56.819857",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -2217,16 +1856,7 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "papermill": {
-     "duration": 0.037029,
-     "end_time": "2020-12-18T15:52:57.371211",
-     "exception": false,
-     "start_time": "2020-12-18T15:52:57.334182",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "source": [
     "The model has fit the linear components of the data well, and also accomodated its heteroskedasticity.  This flexibility, along with the ability to modularly specify the conditional mixture weights and conditional component densities, makes dependent density regression an extremely useful nonparametric Bayesian model.\n",
     "\n",
@@ -2239,43 +1869,28 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2020-12-18T15:53:09.107982Z",
-     "iopub.status.busy": "2020-12-18T15:53:09.106908Z",
-     "iopub.status.idle": "2020-12-18T15:53:09.194464Z",
-     "shell.execute_reply": "2020-12-18T15:53:09.193888Z"
-    },
-    "papermill": {
-     "duration": 0.137799,
-     "end_time": "2020-12-18T15:53:09.194592",
-     "exception": false,
-     "start_time": "2020-12-18T15:53:09.056793",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "execution_count": 17,
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Last updated: Fri Dec 18 2020\n",
+      "Last updated: Sat Apr 24 2021\n",
       "\n",
       "Python implementation: CPython\n",
       "Python version       : 3.7.6\n",
       "IPython version      : 7.13.0\n",
       "\n",
+      "numpy     : 1.17.5\n",
+      "matplotlib: 3.2.1\n",
+      "pymc3     : 3.11.2\n",
+      "arviz     : 0.11.2\n",
+      "theano    : 1.1.2\n",
       "pandas    : 1.1.5\n",
       "seaborn   : 0.10.0\n",
-      "matplotlib: 3.2.1\n",
-      "theano    : 1.0.5\n",
-      "arviz     : 0.10.0\n",
-      "numpy     : 1.17.5\n",
-      "pymc3     : 3.10.0\n",
       "\n",
-      "Watermark: 2.1.0\n",
+      "Watermark: 2.2.0\n",
       "\n"
      ]
     }
@@ -2288,16 +1903,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "papermill": {
-     "duration": 0.042162,
-     "end_time": "2020-12-18T15:53:09.280948",
-     "exception": false,
-     "start_time": "2020-12-18T15:53:09.238786",
-     "status": "completed"
-    },
-    "tags": []
-   },
+   "metadata": {},
    "outputs": [],
    "source": []
   }
@@ -2319,17 +1925,6 @@
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.7.6"
-  },
-  "papermill": {
-   "duration": 844.57815,
-   "end_time": "2020-12-18T15:53:09.535936",
-   "environment_variables": {},
-   "exception": null,
-   "input_path": "__notebook__.ipynb",
-   "output_path": "__notebook__.ipynb",
-   "parameters": {},
-   "start_time": "2020-12-18T15:39:04.957786",
-   "version": "2.1.0"
   }
  },
  "nbformat": 4,