simple Bayesian classifier

bayes-classifier.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "d797047b-992e-41c6-b32d-0dc0a1d86f70",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "%pylab is deprecated, use %matplotlib inline and import the required libraries.\n",
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/duke/env-vbjax/lib/python3.11/site-packages/IPython/core/magics/pylab.py:162: UserWarning: pylab import has clobbered these variables: ['np']\n",
      "`%matplotlib` prevents importing * from pylab and numpy\n",
      "  warn(\"pylab import has clobbered these variables: %s\"  % clobbered +\n"
     ]
    }
   ],
   "source": [
    "%pylab inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "84ebef95-ac80-47b2-a5b0-1810efd483e6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as np\n",
    "import jax.random as jr\n",
    "\n",
    "# make it easier to get a new rng key, just call `key()`\n",
    "keys = [jr.PRNGKey(42)]\n",
    "def key():\n",
    "    keys.append(jr.split(keys[-1],1)[0])\n",
    "    return keys[-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "949f373e-4016-4883-95e2-de0ceafd3cc9",
   "metadata": {},
   "source": [
    "## make model\n",
    "\n",
    "problem dimensions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "e9804eb2-08ed-417a-93a3-6f4164c7773e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "n_class = 5\n",
    "n_obs_dim = 20\n",
    "n_obs = 40\n",
    "confusion = 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a2b727a-5a22-4f49-98e9-b1724ebfcd15",
   "metadata": {},
   "source": [
    "generate random data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "fa2e81f7-627d-4dfb-bee2-20b61109e1ac",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class_means = jr.normal(key(), shape=(n_class, n_obs_dim))\n",
    "obs_classes = jr.randint(key(), shape=(n_obs,), minval=0, maxval=n_class)\n",
    "obs = class_means[obs_classes] + confusion*jr.normal(key(), shape=(n_obs, n_obs_dim))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "950c8276-17a9-427f-b00d-29683f10a460",
   "metadata": {},
   "source": [
    "we want our class ids in \"one hot\" encoding:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "cd81650f-05b6-49d6-9de1-62491d923fef",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([4, 0, 3, 0, 2], dtype=int32),\n",
       " Array([[0., 1., 0., 1., 0.],\n",
       "        [0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 1.],\n",
       "        [0., 0., 1., 0., 0.],\n",
       "        [1., 0., 0., 0., 0.]], dtype=float32))"
      ]
     },
     "execution_count": 83,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "obs_classes_oh = jax.nn.one_hot(obs_classes, num_classes=n_class).T\n",
    "obs_classes[:5], obs_classes_oh[:,:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee1bc227-c4d9-4445-9fc7-84893d81068c",
   "metadata": {},
   "source": [
    "a forward model for class probabilities. first, the parameters,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "6269dece-74e1-443d-88f9-db56c3f586b3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mix = jr.normal(key(), shape=(n_class, n_obs_dim))\n",
    "offset = jr.normal(key(), shape=(n_class,1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e612d634-61d3-4d1b-92c9-7ed8e4ca1452",
   "metadata": {},
   "source": [
    "then a model which predicts class likelihoods with a linear model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "76b82429-66a8-4200-86b1-79ecd8b84e59",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 5, 60,  0, 36, 35],\n",
       "       [ 0,  0,  0,  0,  0],\n",
       "       [ 0,  2,  5,  1, 61],\n",
       "       [ 0, 37, 94, 62,  1],\n",
       "       [93,  0,  0,  0,  0]], dtype=int32)"
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def fwd(params):\n",
    "    mix, offset = params\n",
    "    # use linear model to predict class based on observations\n",
    "    linear_pred = mix@obs.T + offset\n",
    "    # normalize them, \"one-hot\"\n",
    "    norm_pred = jax.nn.softmax(linear_pred, axis=0)\n",
    "    return norm_pred\n",
    " \n",
    "(fwd((mix, offset))*100).astype('i')[:,:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18823d3c-9efe-4994-8f09-dd123fae0a00",
   "metadata": {},
   "source": [
    "each column has the percentage likelihood of the class (each class per row) for that observation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82701ba0-fc73-4e88-94eb-016973b3c448",
   "metadata": {},
   "source": [
    "## optimize\n",
    "\n",
    "now a loss function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "82f29c9f-8d91-4539-86ab-b98c9ac481a9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 50.402283\n",
      "100 1.8094155\n",
      "200 0.9590366\n",
      "300 0.65946054\n",
      "400 0.5048584\n",
      "500 0.4100145\n",
      "600 0.34570068\n",
      "700 0.29913133\n",
      "800 0.26380706\n",
      "900 0.23606819\n"
     ]
    }
   ],
   "source": [
    "def loss(params):\n",
    "    mix, offset = params\n",
    "    p = fwd(params)\n",
    "    lp = obs_classes_oh*np.log(p) + (1-obs_classes_oh)*np.log(1-p)\n",
    "    return -np.sum(lp)\n",
    "\n",
    "grad_loss = jax.jit(jax.grad(loss))\n",
    "\n",
    "params = mix, offset\n",
    "# loss(params), grad_loss(params)\n",
    "for i in range(1000):\n",
    "    gmix, goffset = grad_loss(params)\n",
    "    mix = mix - 1e-3*gmix\n",
    "    offset = offset - 1e-3*goffset\n",
    "    params = mix, offset\n",
    "    if i%100 == 0:\n",
    "        print(i, loss(params))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67c0a6fd-a6d3-4959-bcfd-fcc4ff3d4624",
   "metadata": {},
   "source": [
    "Now compare predictions and true labels,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "65b915b2-a3cc-4970-a14a-e57576983e08",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([4, 0, 3, 0, 2], dtype=int32), Array([4, 0, 3, 0, 2], dtype=int32))"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.argmax((fwd((mix, offset))*100).astype('i')[:,:5], axis=0), obs_classes[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81095783-0f58-4ef1-94e0-60a6e64db026",
   "metadata": {},
   "source": [
    "More tricks explained in this page https://www.architecture-performance.fr/ap_blog/logistic-regression-with-jax/\n",
    "\n",
    "## MCMC\n",
    "\n",
    "but we can try to do Bayesian inference with the model already"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ce3247d-dbb6-49e8-a9ad-5f988d89411c",
   "metadata": {},
   "source": [
    "What's the equivalent of our loss function?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "id": "2dd52503-53fa-4b9d-9b43-7f1b1a514d04",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([4, 0, 3, 0, 2], dtype=int32)"
      ]
     },
     "execution_count": 132,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpyro, numpyro.infer\n",
    "from numpyro.distributions import Normal, Bernoulli\n",
    "\n",
    "b = Bernoulli(probs=fwd((mix, offset)))\n",
    "np.argmax(-b.log_prob(obs_classes_oh)[:,:5], axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e50588f9-5ca2-4ac2-b2a7-79275bfd690b",
   "metadata": {},
   "source": [
    "so now a log probability function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "id": "64ec1cd6-76c6-4fc7-b044-e3056f2d7561",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def log_p():\n",
    "    mix_mu = np.zeros((n_class, n_obs_dim))\n",
    "    offset_mu = np.zeros((n_class,1))\n",
    "    mix = numpyro.sample('mix', Normal(mix_mu, 1))\n",
    "    offset = numpyro.sample('offset', Normal(offset_mu, 1))\n",
    "    params = mix, offset\n",
    "    p = fwd(params)\n",
    "    numpyro.sample('obs_classes_oh',\n",
    "                   Bernoulli(probs=p),\n",
    "                   obs=obs_classes_oh)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3a0fc31-2eb1-4102-bdab-2fc8c92fde67",
   "metadata": {},
   "source": [
    "now run it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "id": "080ccd9d-f4c3-437e-97a4-53974d0548ac",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample: 100%|███████████████████████████████████████| 1000/1000 [00:01<00:00, 891.13it/s, 31 steps of size 1.28e-01. acc. prob=0.93]\n"
     ]
    }
   ],
   "source": [
    "mcmc = numpyro.infer.MCMC(\n",
    "    numpyro.infer.NUTS(log_p),\n",
    "    num_warmup=500,\n",
    "    num_samples=500\n",
    ")\n",
    "mcmc_key = jax.random.PRNGKey(42)\n",
    "mcmc.run(mcmc_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "id": "954e5c2f-cbc3-4871-820f-2c2c7b05000d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mix_h = mcmc.get_samples()['mix']\n",
    "offset_h = mcmc.get_samples()['offset']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7061dbe9-3738-4280-ac66-c21cf8308168",
   "metadata": {},
   "source": [
    "Let's first check if the chain mixes well?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "id": "0ffb7c3e-8a4d-4488-9714-20dddfd0401f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x2c6961450>"
      ]
     },
     "execution_count": 169,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAvcAAADICAYAAABszM7qAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAucUlEQVR4nO3dfVQUV54//nd3K40PNIYYQQZQEvOgUTHRgZDMzoAyAuEY3fE30Ww2wceMDiTRnqzKWRVjdkI2TnwYZTXfUSSexESzG8wkeDCIQcYJPoFMfJh4YoYoKg3mQVoQeei6vz8cKim7GroUqJJ+v86559i3bt3P7bJoP15u3zIJIQSIiIiIiOi2Z9Z7AERERERE1DmY3BMRERER9RBM7omIiIiIeggm90REREREPQSTeyIiIiKiHoLJPRERERFRD8HknoiIiIioh2ByT0RERETUQzC5JyIiIiLqIZjcExERERH1EL26O6AkSbh48SICAgJgMpm6OzwRERER3UAIgStXriA0NBRms7Hmfq9du4bm5maPx/38/ODv7+9VXyUlJVi1ahXKyspQXV2NvLw8TJkypd1ziouLYbfbcfLkSYSHh2Pp0qWYMWOGpj6FEMjMzMSf/vQnXL58GY899hg2btyIe++9V27z3Xff4fnnn8dHH30Es9mMqVOnYt26dejfv79X7+3HwbpVVVWVAMDCwsLCwsLCwmKwUlVV1d2pYbsaGxtFyCBLu2MOCQkRjY2NXvW3e/du8Z//+Z/igw8+EABEXl5eu+3/8Y9/iL59+wq73S5OnTol1q9fLywWiygoKNDU52uvvSYCAwPFrl27xN/+9jfxxBNPiMjISMW4k5KSRFRUlDh48KD4y1/+IoYNGyaeeuopr97Xj5mEEALdqK6uDgMGDMDPkIJept7dGRoAYLb6dXvMNhfnjNEt9k92/kOXuKK+Xpe4ACBckm6xMeJu3UKbvzqvW+yrsfd23KiLfH9v93+eAEBI9mFd4gKAyWLRLXZT3GjdYvvt+5tusSF0/Fzp3n+uFcwj7tMncNVFfeICgMulW2jR3NLtMVtFC/7i+jMuX76MwMDAbo/vidPpRGBgIM4cDYctwP03Cs4rEoaNq0JdXR1sNpumvk0mU4cz94sXL0Z+fj5OnDgh102fPh2XL19GQUGBV30KIRAaGorf/e53eOmllwBcz4eDg4ORm5uL6dOn4+9//ztGjBiBI0eOYNy4cQCAgoICPP744zh//jxCQ0O9fl/dviynbSlOL1NvfZJ7k37JvcXq3a+MukIvsz7vW+h4vYVJvw9mWKy6hdbzHu/VW7973GLVJ7nX43OsjcmkX3Lv0vHvWs9rDuiY3EPH5F6vzzQdP8+g478hQsdVy0ZdMt0/wIT+Ae5jk3C9zul0KuqtVius1lu/b0tLS5GQkKCoS0xMxIIFC7zuo7KyEg6HQ9FPYGAgYmJiUFpaiunTp6O0tBQDBgyQE3sASEhIgNlsxqFDh/Cv//qvXscz1qIqIiIiIqIbtAiXxwIA4eHhCAwMlEtWVlanxHU4HAgODlbUBQcHw+l0orGx0es+2s67sZ+2Yw6HA4MGDVIc79WrF4KCguQ23ur2mXsiIiIiIi1aIUFtsVLrP3+bVlVVpViW0xmz9rcrztwTERERkaFJEB4LANhsNkXprOQ+JCQENTU1irqamhrYbDb06dPH6z7azruxn7ZjISEhqK2tVRxvbW3Fd999J7fxFpN7IiIiIjK0FiE8lq4UGxuLoqIiRV1hYSFiY2O97iMyMhIhISGKfpxOJw4dOiT3Exsbi8uXL6OsrExus2/fPkiShJiYGE1j5rIcIiIiIjK0ZiHQrJLIq9W1p76+HmfOnJFfV1ZWoqKiAkFBQYiIiEBGRgYuXLiAbdu2AQDmzZuHDRs2YNGiRZg1axb27duHnTt3Ij8/3+s+TSYTFixYgP/6r//Cvffei8jISCxbtgyhoaHyrjrDhw9HUlIS5s6di02bNqGlpQXp6emYPn26pp1yACb3RERERGRwEtT3qtK6f9XRo0cRHx8vv7bb7QCA1NRU5Obmorq6GufOnZOPR0ZGIj8/HwsXLsS6desQFhaGzZs3IzEx0es+AWDRokVoaGjAc889h8uXL+NnP/sZCgoKFA/feuedd5Ceno4JEybID7H64x//qPEdMrknIiIiIoNrFSa0qOwR2qpx39C4uDi094intmT8xnOOHTt2030C17cYXblyJVauXOmxTVBQELZv395uP95gck9EREREhuaCCS64J/Jqdb6OyT0RERERGVqLMKNFuO8D06Lfs90M66Z2y8nOzsbQoUPh7++PmJgYHD6s3+PXiYiIiKhna4bFYyElzcn9jh07YLfbkZmZifLyckRFRSExMdFtb04iIiIios4ghAmSShEa19z7As3J/erVqzF37lzMnDkTI0aMwKZNm9C3b1/k5OR0xfiIiIiIyMc1C4vHQkqakvvm5maUlZUhISHhhw7MZiQkJKC0tFT1nKamJjidTkUhIiIiIvJWC8xogUWl8HmsN9J0Rb755hu4XC4EBwcr6oODg+FwOFTPycrKQmBgoFzCw8NvfrRERERE5HNcwuyxkFKXX5GMjAzU1dXJpaqqqqtDEhEREVEP0qo6a29BK79Q60bTVpgDBw6ExWJBTU2Nor6mpgYhISGq51itVlit1psfIRERERH5tBbRCy0q6+vVHmzl6zTN3Pv5+WHs2LEoKiqS6yRJQlFREWJjYzt9cERERERELmHyWEhJ80Os7HY7UlNTMW7cOERHR2Pt2rVoaGjAzJkzu2J8REREROTjOHPvPc3J/bRp03Dp0iUsX74cDocDY8aMQUFBgduXbImIiIiIOkMLzKrbXraAj6i9kebkHgDS09ORnp7e2WMhIiIiInIjwQxJZTW5Wp2vu6nknoiIiIiou7QIC3qpLsvhzP2NmNwTERERkaE1i16wCPe0tZm5vRsm90RERERkaJIwQVL58qxana9jck9EREREhtYqeqFFZea+lTP3bpjcExEREZGhtQgLLFxz7xUm90RERERkaC4ALrgvwXF1/1AMj8k9ERERERlai9QLFsk9bW2ROHN/Iyb3RERERGRorcKi+oTaViHpMBpj0y25twTdAYvZr9vjmnr37vaYbQIrdfzlUWurLmGlq1d1iQsA5r59dYv9/fAA3WJHrNPvHseUf+gW2pr/rS5xLQMCdYkLACYdY1uu6fd5VjsvRrfYgzaW6hZbTxd/GaRL3ND/V6VLXAAQzc36xdZhNloYfO16Z+2WU1JSglWrVqGsrAzV1dXIy8vDlClT2j2nuLgYdrsdJ0+eRHh4OJYuXYoZM2Yo2mRnZ2PVqlVwOByIiorC+vXrER0dDQD4+uuvERkZqdr3zp078etf/xoAYDK5v5d3330X06dP1/Qe+VgvIiIiIjK0ln/O3KsVLRoaGhAVFYXs7Gyv2ldWViIlJQXx8fGoqKjAggULMGfOHOzZs0dus2PHDtjtdmRmZqK8vBxRUVFITExEbW0tACA8PBzV1dWK8vLLL6N///5ITk5WxNu6dauiXUf/8VDDZTlEREREZGgtwgKz6m4515flOJ1ORb3VaoXVanVrn5yc7JZQt2fTpk2IjIzEG2+8AQAYPnw4Dhw4gDVr1iAxMREAsHr1asydOxczZ86Uz8nPz0dOTg6WLFkCi8WCkJAQRb95eXl48skn0b9/f0X9gAED3NpqxZl7IiIiIjK0tmU5agW4PjseGBgol6ysrE6JW1paioSEBEVdYmIiSkuvL9Frbm5GWVmZoo3ZbEZCQoLc5kZlZWWoqKjA7Nmz3Y6lpaVh4MCBiI6ORk5Ozk0tl+LMPREREREZWquHmfu2L9RWVVXBZrPJ9Wqz9jfD4XAgODhYURccHAyn04nGxkZ8//33cLlcqm2++OIL1T63bNmC4cOH49FHH1XUr1y5EuPHj0ffvn3xySef4Le//S3q6+vxwgsvaBozk3siIiIiMrRWyQyzpJLcS9e/3G+z2RTJvVE1NjZi+/btWLZsmduxH9c99NBDaGhowKpVqzQn91yWQ0RERESGJsHksXSlkJAQ1NTUKOpqampgs9nQp08fDBw4EBaLRbWN2tr5//3f/8XVq1fx7LPPdhg7JiYG58+fR1NTk6YxM7knIiIiIkNrkSweS1eKjY1FUVGRoq6wsBCxsbEAAD8/P4wdO1bRRpIkFBUVyW1+bMuWLXjiiSdw1113dRi7oqICd9xxh+YlRlyWQ0RERESG5oIZrcJ9TtqlcZ66vr4eZ86ckV9XVlaioqICQUFBiIiIQEZGBi5cuIBt27YBAObNm4cNGzZg0aJFmDVrFvbt24edO3ciPz9f7sNutyM1NRXjxo1DdHQ01q5di4aGBnn3nDZnzpxBSUkJdu/e7Taujz76CDU1NXjkkUfg7++PwsJCvPrqq3jppZc0vT+AyT0RERERGVxnPcTq6NGjiI+Pl1/b7XYAQGpqKnJzc1FdXY1z587JxyMjI5Gfn4+FCxdi3bp1CAsLw+bNm+VtMAFg2rRpuHTpEpYvXw6Hw4ExY8agoKDA7Uu2OTk5CAsLw8SJE93G1bt3b2RnZ2PhwoUQQmDYsGHyFptaMbknIiIiIkNrlSwwqX6hVtuynLi4uHa3l8zNzVU959ixY+32m56ejvT09HbbvPrqq3j11VdVjyUlJSEpKand873F5J6IiIiIDK1VmGFSWZajtlTH12m+IiUlJZg0aRJCQ0NhMpmwa9euLhgWEREREdF1HT3Ein6gOblvaGhAVFQUsrOzu2I8REREREQKrZLZYyElzctykpOTkZyc7HX7pqYmxf6cTqdTa0giIiIi8mEuYVJdluPizL2bLv/vTlZWFgIDA+USHh7e1SGJiIiIqAfhshzvdXlyn5GRgbq6OrlUVVV1dUgiIiIi6kG4LMd7Xb5bjtVq1fxkLSIiIiKiNi4Pu+W4uFuOG26FSURERESG1lkPsfIFTO6JiIiIyNBckhkmlSU4Li7LcaM5ua+vr8eZM2fk15WVlaioqEBQUBAiIiI6dXBERERERJJkVk3kJSb3bjQn90ePHkV8fLz82m63AwBSU1NVH9lLRERERHQrBAAh1OtJSXNyHxcXB6F2dYmIiIiIuoBLmAF+odYrXHNPRERERIbmkkyA5P7lWZdKna9jck9EREREhiaECUJlZxy1Ol/H5J6IiIiIDM0lmQHuluMVJvdEREREZGiSBJhUluBIkg6DMTgm90RERERkaFyW4z0m90RERERkaJIwwcQn1HqFyT0RERERGZtkglDbGYe75bhhck9EREREhiaEh4dY8dFLbnRL7psfDIfUy7/b4174RffHbBOx4jPdYouAAH0Cm/T7H7VobtYtdtCfT+kW+0LTcN1iD+h9VrfYZn99frZdl+t0iQsATY/dp1vsficcusUeVFylW2w9ST8bo1vskDX6/Pv19YpHdYkLABEvl+oWu9eQ8O4PKjUB+n2Ed0hIZgiVnXHU6tpTUlKCVatWoaysDNXV1cjLy8OUKVPaPae4uBh2ux0nT55EeHg4li5dihkzZijaZGdnY9WqVXA4HIiKisL69esRHR0tH4+Li8P+/fsV5/zmN7/Bpk2b5Nfnzp3D/Pnz8emnn6J///5ITU1FVlYWevXSlq5z/yAiIiIiMjQheS5aNDQ0ICoqCtnZ2V61r6ysREpKCuLj41FRUYEFCxZgzpw52LNnj9xmx44dsNvtyMzMRHl5OaKiopCYmIja2lpFX3PnzkV1dbVcXn/9dfmYy+VCSkoKmpub8dlnn+Gtt95Cbm4uli9fru0NgstyiIiIiMjgOmu3nOTkZCQnJ3vdftOmTYiMjMQbb7wBABg+fDgOHDiANWvWIDExEQCwevVqzJ07FzNnzpTPyc/PR05ODpYsWSL31bdvX4SEhKjG+eSTT3Dq1Cns3bsXwcHBGDNmDF555RUsXrwYK1asgJ+fn9dj5sw9ERERERmaENe/UOtW/pncO51ORWlqauqUuKWlpUhISFDUJSYmorT0+rKt5uZmlJWVKdqYzWYkJCTIbdq88847GDhwIEaOHImMjAxcvXpVEWfUqFEIDg5WxHE6nTh58qSmMXPmnoiIiIiMTZiuF7V6AOHhyu8pZGZmYsWKFbcc1uFwKBJuAAgODobT6URjYyO+//57uFwu1TZffPGF/Prf/u3fMGTIEISGhuLzzz/H4sWLcfr0aXzwwQftxmk7pgWTeyIiIiIyNvHPolYPoKqqCjabTa62Wq3dMixvPffcc/KfR40ahcGDB2PChAn46quvcM8993RqLC7LISIiIiJjk0yeCwCbzaYonZXch4SEoKamRlFXU1MDm82GPn36YODAgbBYLKptPK2vB4CYmBgAwJkzZ9qN03ZMCyb3RERERGRonbVbjlaxsbEoKipS1BUWFiI2NhYA4Ofnh7FjxyraSJKEoqIiuY2aiooKAMDgwYPlOMePH1fssFNYWAibzYYRI0ZoGjOX5RARERGRsXWw5t5b9fX18mw5cH2ry4qKCgQFBSEiIgIZGRm4cOECtm3bBgCYN28eNmzYgEWLFmHWrFnYt28fdu7cifz8fLkPu92O1NRUjBs3DtHR0Vi7di0aGhrk3XO++uorbN++HY8//jjuvPNOfP7551i4cCF+/vOfY/To0QCAiRMnYsSIEXjmmWfw+uuvw+FwYOnSpUhLS9P8Wwgm90RERERkaCbpelGr1+Lo0aOIj4+XX9vtdgBAamoqcnNzUV1djXPnzsnHIyMjkZ+fj4ULF2LdunUICwvD5s2b5W0wAWDatGm4dOkSli9fDofDgTFjxqCgoED+Qqyfnx/27t0rJ/3h4eGYOnUqli5dKvdhsVjw8ccfY/78+YiNjUW/fv2QmpqKlStXanuDYHJPREREREb3o/X1bvUaxMXFQQi1b+Zel5ubq3rOsWPH2u03PT0d6enpqsfCw8Pdnk6rZsiQIdi9e3eH7Tqiac19VlYWfvrTnyIgIACDBg3ClClTcPr06VseBBERERGRR6KdQgqakvv9+/cjLS0NBw8eRGFhIVpaWjBx4kQ0NDR01fiIiIiIyNd1sFsO/UDTspyCggLF69zcXAwaNAhlZWX4+c9/3qkDIyIiIiICOm/NvS+4pTX3dXV1AICgoCCPbZqamhSPAHY6nbcSkoiIiIiIPLjpfe4lScKCBQvw2GOPYeTIkR7bZWVlITAwUC43Ph6YiIiIiKg9JmGCSVIpGrfC9AU3ndynpaXhxIkTeO+999ptl5GRgbq6OrlUVVXdbEgiIiIi8kVSO4UUbmpZTnp6Oj7++GOUlJQgLCys3bZWq7XTHgFMRERERL7HJK4XtXpS0pTcCyHw/PPPIy8vD8XFxYiMjOyqcRERERERXedplp4z9240JfdpaWnYvn07PvzwQwQEBMDhcAAAAgMD0adPny4ZIBERERH5trY19mr1pKRpzf3GjRtRV1eHuLg4DB48WC47duzoqvERERERka/jQ6y8pnlZDhERERFRd+I+9967pX3uiYiIiIi6nIfknmvu3TG5JyIiIiJj87QEh4tK3DC5JyIiIiJD41aY3mNyT0RERETGxpl7rzG5JyIiIiJDMwkPX6hlcu+GyT0RERERGRp3y/Eek3siIiIiMjYuy/Eak3siIiIiMjTO3HtPt+TeevZb9DJbuz1u5InGbo/ZRuql3/+lpCtXdIutF1OfPrrFlu6L0C22Lf+4brFFX/2uuSkgQJ+4Lv3+ZXGG6/eZ4r+3VrfY5r59dYstXWvSLXZr/966xfYzW3SJO/f/K9AlLgDsWWHTLXbzkDu7PWZr6zXgbLeH9RqTe+9x5p6IiIiIjI3Lcrxm1nsARERERETtaZu5VytalJSUYNKkSQgNDYXJZMKuXbs6PKe4uBgPP/wwrFYrhg0bhtzcXLc22dnZGDp0KPz9/RETE4PDhw/Lx7777js8//zzuP/++9GnTx9ERETghRdeQF1dnfI9mkxu5b333tP2BsHknoiIiIgMrrOS+4aGBkRFRSE7O9ur9pWVlUhJSUF8fDwqKiqwYMECzJkzB3v27JHb7NixA3a7HZmZmSgvL0dUVBQSExNRW3t9+eLFixdx8eJF/OEPf8CJEyeQm5uLgoICzJ492y3e1q1bUV1dLZcpU6Zoe4PgshwiIiIiMroOluU4nU5FtdVqhdXq/t3O5ORkJCcnex1206ZNiIyMxBtvvAEAGD58OA4cOIA1a9YgMTERALB69WrMnTsXM2fOlM/Jz89HTk4OlixZgpEjR+L//u//5D7vuece/P73v8e///u/o7W1Fb1+9J3MAQMGICQkxOvxqeHMPREREREZmkl4LgAQHh6OwMBAuWRlZXVK3NLSUiQkJCjqEhMTUVpaCgBobm5GWVmZoo3ZbEZCQoLcRk1dXR1sNpsisQeAtLQ0DBw4ENHR0cjJyYEQ2r9UwJl7IiIiIjK0jnbLqaqqgs32ww5HarP2N8PhcCA4OFhRFxwcDKfTicbGRnz//fdwuVyqbb744gvVPr/55hu88soreO655xT1K1euxPjx49G3b1988skn+O1vf4v6+nq88MILmsbM5J6IiIiIjK2DZTk2m02R3BuV0+lESkoKRowYgRUrViiOLVu2TP7zQw89hIaGBqxatUpzcs9lOURERERkaCbh4Qu1XbwVZkhICGpqahR1NTU1sNls6NOnDwYOHAiLxaLa5sa181euXEFSUhICAgKQl5eH3r3bf3ZFTEwMzp8/j6Ymbc/XYHJPRERERIbWWbvlaBUbG4uioiJFXWFhIWJjYwEAfn5+GDt2rKKNJEkoKiqS2wDXZ+wnTpwIPz8//PnPf4a/v3+HsSsqKnDHHXdoXmLEZTlEREREZGyd9BCr+vp6nDlzRn5dWVmJiooKBAUFISIiAhkZGbhw4QK2bdsGAJg3bx42bNiARYsWYdasWdi3bx927tyJ/Px8uQ+73Y7U1FSMGzcO0dHRWLt2LRoaGuTdc9oS+6tXr+Ltt9+G0+mUd/e56667YLFY8NFHH6GmpgaPPPII/P39UVhYiFdffRUvvfSStjcIJvdEREREZHAdfaHWW0ePHkV8fLz82m63AwBSU1ORm5uL6upqnDt3Tj4eGRmJ/Px8LFy4EOvWrUNYWBg2b94sb4MJANOmTcOlS5ewfPlyOBwOjBkzBgUFBfKXbMvLy3Ho0CEAwLBhwxTjqaysxNChQ9G7d29kZ2dj4cKFEEJg2LBh8habWmlK7jdu3IiNGzfi66+/BgA8+OCDWL58uab9QomIiIiItOis5D4uLq7d7SXVnj4bFxeHY8eOtdtveno60tPTbyomACQlJSEpKandNt7StOY+LCwMr732GsrKynD06FGMHz8ekydPxsmTJztlMEREREREbkQ7hRQ0zdxPmjRJ8fr3v/89Nm7ciIMHD+LBBx/s1IEREREREQGASRIwSe6ZvFqdr7vpNfculwvvv/8+GhoaFN8GvlFTU5NiC58bHw9MRERERNSezlqW4ws0b4V5/Phx9O/fH1arFfPmzUNeXh5GjBjhsX1WVpbiccDh4eG3NGAiIiIi8i0m4bmQkubk/v7770dFRQUOHTqE+fPnIzU1FadOnfLYPiMjA3V1dXKpqqq6pQETERERkW/Ra5/725HmZTl+fn7yNj5jx47FkSNHsG7dOrz55puq7a1Wq+bN94mIiIiIZMJDIs+Zeze3vM+9JEmaH4tLREREROQ1Ia4XtXpS0JTcZ2RkIDk5GREREbhy5Qq2b9+O4uJi7Nmzp6vGR0REREQ+jl+o9Z6m5L62thbPPvssqqurERgYiNGjR2PPnj345S9/2VXjIyIiIiIfZ3IBJpVvippc3T8Wo9OU3G/ZsqWrxkFEREREpMrTzjjcLcfdLa+5JyIiIiLqSnyIlfeY3BMRERGRoXHNvfeY3BMRERGRsXG3HK8xuSciIiIiQ+PMvfeY3BMRERGRoZlcAiazypp7F2fub8TknoiIiIiMTUD9abTM7d0wuSciIiIiQzMJD7vlcM29Gyb3RERERGRoXHPvPd2S+6+n/QQWq3+3xw09cK3bY7bxO6HjHdjaqktY0dyiS1wAkOrrdYv9jycDdIt99+JG3WLj6lXdQjdMjdElbuBfv9YlLgDctalUt9jo7adb6MaEB3WL3fdrp26xW/uqPJ6zm/hJ+jwG1B70D13iAsAe81jdYkOPdeQGX7tuEkJ1lp4z9+44c09EREREhmZyCZhUHkfLL9S6028agIiIiIjIG5LwXDQoKSnBpEmTEBoaCpPJhF27dnV4TnFxMR5++GFYrVYMGzYMubm5bm2ys7MxdOhQ+Pv7IyYmBocPH1Ycv3btGtLS0nDnnXeif//+mDp1KmpqahRtzp07h5SUFPTt2xeDBg3Cf/zHf6D1JlZeMLknIiIiIkMzCc9Fi4aGBkRFRSE7O9ur9pWVlUhJSUF8fDwqKiqwYMECzJkzB3v27JHb7NixA3a7HZmZmSgvL0dUVBQSExNRW1srt1m4cCE++ugjvP/++9i/fz8uXryIX/3qV/Jxl8uFlJQUNDc347PPPsNbb72F3NxcLF++XNsbBJflEBEREZHBddaynOTkZCQnJ3vdftOmTYiMjMQbb7wBABg+fDgOHDiANWvWIDExEQCwevVqzJ07FzNnzpTPyc/PR05ODpYsWYK6ujps2bIF27dvx/jx4wEAW7duxfDhw3Hw4EE88sgj+OSTT3Dq1Cns3bsXwcHBGDNmDF555RUsXrwYK1asgJ+f999z4sw9ERERERlbB8tynE6nojQ1NXVK2NLSUiQkJCjqEhMTUVp6fUOD5uZmlJWVKdqYzWYkJCTIbcrKytDS0qJo88ADDyAiIkJuU1pailGjRiE4OFgRx+l04uTJk5rGzOSeiIiIiAytbbcctQIA4eHhCAwMlEtWVlanxHU4HIqEGwCCg4PhdDrR2NiIb775Bi6XS7WNw+GQ+/Dz88OAAQPabaPWR9sxLbgsh4iIiIiMTRLq23X+c+a+qqoKNptNrrZard01MsNhck9EREREhmaSBEwqT6xqe2qtzWZTJPedJSQkxG1Xm5qaGthsNvTp0wcWiwUWi0W1TUhIiNxHc3MzLl++rJi9v7HNjTvstPXZ1sZbXJZDRERERMYmhOfShWJjY1FUVKSoKywsRGxsLADAz88PY8eOVbSRJAlFRUVym7Fjx6J3796KNqdPn8a5c+fkNrGxsTh+/Lhih53CwkLYbDaMGDFC05g5c09EREREhmZyCZhw67vl1NfX48yZM/LryspKVFRUICgoCBEREcjIyMCFCxewbds2AMC8efOwYcMGLFq0CLNmzcK+ffuwc+dO5Ofny33Y7XakpqZi3LhxiI6Oxtq1a9HQ0CDvnhMYGIjZs2fDbrcjKCgINpsNzz//PGJjY/HII48AACZOnIgRI0bgmWeeweuvvw6Hw4GlS5ciLS1N8xIjJvdEREREZGySBKgsy4GkUteOo0ePIj4+Xn5tt9sBAKmpqcjNzUV1dTXOnTsnH4+MjER+fj4WLlyIdevWISwsDJs3b5a3wQSAadOm4dKlS1i+fDkcDgfGjBmDgoICxRdk16xZA7PZjKlTp6KpqQmJiYn4n//5H/m4xWLBxx9/jPnz5yM2Nhb9+vVDamoqVq5cqen9AbeY3L/22mvIyMjAiy++iLVr195KV0RERERE6jwtwdG4LCcuLg6inXPUnj4bFxeHY8eOtdtveno60tPTPR739/dHdnZ2uw/PGjJkCHbv3t1uHG/cdHJ/5MgRvPnmmxg9evQtD4KIiIiIyJPOWpbjC27qC7X19fV4+umn8ac//Ql33HFHZ4+JiIiIiOgHLslzIYWbSu7T0tKQkpLi9sQuNU1NTW5PDSMiIiIi8ppOu+XcjjQvy3nvvfdQXl6OI0eOeNU+KysLL7/8suaBEREREREBAISk/uVZwZn7G2maua+qqsKLL76Id955B/7+/l6dk5GRgbq6OrlUVVXd1ECJiIiIyEe5XJ4LKWiauS8rK0NtbS0efvhhuc7lcqGkpAQbNmxAU1MTLBaL4hyr1erTjwAmIiIiolvUSbvl+AJNyf2ECRNw/PhxRd3MmTPxwAMPYPHixW6JPRERERHRLXNJ6ktwNO5z7ws0JfcBAQEYOXKkoq5fv36488473eqJiIiIiDqFJACoJfecub8Rn1BLRERERMYmSVBP7jlzf6NbTu6Li4s7YRhERERERB4wufcaZ+6JiIiIyNCEywUh3HfGERJ3y7kRk3siIiIiMjYh1NfXc7ccN0zuiYiIiMjYXC7ApDJLrzKb7+uY3BMRERGRoQmXC0IluVdbquPrmNwTERERkbEJAYDLcrzB5J6IiIiIjM0leViWw91ybsTknoiIiIgM7fqyHLN7PZfluOn25F7889cnUtO17g4NAGht1ScuAJilZt1iQ7TqE1a06BIXACQ9Y1/T7z5r1fF96/nr0dYWnT5TdPy51vPv2iRMusXW6+8aAFpdTfrF1vN963SvOa/oNyur58+XpEOu0tp6/d4WBl3m0iI1Q6gsy2mFjv/mGZRJdPPf4vnz5xEeHt6dIYmIiIjIC1VVVQgLC9N7GLJr164hMjISDofDY5uQkBBUVlbC39+/G0dmXN2e3EuShIsXLyIgIAAmk7aZH6fTifDwcFRVVcFms3XRCHsWXjPteM204zXTjtdMO14z7XjNtPPVayaEwJUrVxAaGgqz2X35i56uXbuG5mbPvyX18/NjYv8j3b4sx2w23/L/CG02m0/9wHUGXjPteM204zXTjtdMO14z7XjNtPPFaxYYGKj3EFT5+/szedfAWP81IyIiIiKim8bknoiIiIioh7itknur1YrMzExYrVa9h3Lb4DXTjtdMO14z7XjNtOM1047XTDteM7rddfsXaomIiIiIqGvcVjP3RERERETkGZN7IiIiIqIegsk9EREREVEPweSeiIiIiKiHYHJPRERERNRDGC65z87OxtChQ+Hv74+YmBgcPny43fbvv/8+HnjgAfj7+2PUqFHYvXt3N41Uf1lZWfjpT3+KgIAADBo0CFOmTMHp06fbPSc3Nxcmk0lRfOmpbytWrHB7/w888EC75/jyPQYAQ4cOdbtmJpMJaWlpqu198R4rKSnBpEmTEBoaCpPJhF27dimOCyGwfPlyDB48GH369EFCQgK+/PLLDvvV+nl4O2nvmrW0tGDx4sUYNWoU+vXrh9DQUDz77LO4ePFiu33ezM/37aSj+2zGjBlu7z8pKanDfn31PgOg+tlmMpmwatUqj3329PuMbn+GSu537NgBu92OzMxMlJeXIyoqComJiaitrVVt/9lnn+Gpp57C7NmzcezYMUyZMgVTpkzBiRMnunnk+ti/fz/S0tJw8OBBFBYWoqWlBRMnTkRDQ0O759lsNlRXV8vl7Nmz3TRiY3jwwQcV7//AgQMe2/r6PQYAR44cUVyvwsJCAMCvf/1rj+f42j3W0NCAqKgoZGdnqx5//fXX8cc//hGbNm3CoUOH0K9fPyQmJuLatWse+9T6eXi7ae+aXb16FeXl5Vi2bBnKy8vxwQcf4PTp03jiiSc67FfLz/ftpqP7DACSkpIU7//dd99tt09fvs8AKK5VdXU1cnJyYDKZMHXq1Hb77cn3GfUAwkCio6NFWlqa/NrlconQ0FCRlZWl2v7JJ58UKSkpirqYmBjxm9/8pkvHaVS1tbUCgNi/f7/HNlu3bhWBgYHdNyiDyczMFFFRUV635z3m7sUXXxT33HOPkCRJ9biv32MARF5envxakiQREhIiVq1aJdddvnxZWK1W8e6773rsR+vn4e3sxmum5vDhwwKAOHv2rMc2Wn++b2dq1yw1NVVMnjxZUz+8z5QmT54sxo8f324bX7rP6PZkmJn75uZmlJWVISEhQa4zm81ISEhAaWmp6jmlpaWK9gCQmJjosX1PV1dXBwAICgpqt119fT2GDBmC8PBwTJ48GSdPnuyO4RnGl19+idDQUNx99914+umnce7cOY9teY8pNTc34+2338asWbNgMpk8tvP1e+zHKisr4XA4FPdRYGAgYmJiPN5HN/N52NPV1dXBZDJhwIAB7bbT8vPdExUXF2PQoEG4//77MX/+fHz77bce2/I+U6qpqUF+fj5mz57dYVtfv8/I2AyT3H/zzTdwuVwIDg5W1AcHB8PhcKie43A4NLXvySRJwoIFC/DYY49h5MiRHtvdf//9yMnJwYcffoi3334bkiTh0Ucfxfnz57txtPqJiYlBbm4uCgoKsHHjRlRWVuJf/uVfcOXKFdX2vMeUdu3ahcuXL2PGjBke2/j6PXajtntFy310M5+HPdm1a9ewePFiPPXUU7DZbB7baf357mmSkpKwbds2FBUV4b//+7+xf/9+JCcnw+Vyqbbnfab01ltvISAgAL/61a/abefr9xkZXy+9B0CdIy0tDSdOnOhw3V9sbCxiY2Pl148++iiGDx+ON998E6+88kpXD1N3ycnJ8p9Hjx6NmJgYDBkyBDt37vRqtsbXbdmyBcnJyQgNDfXYxtfvMepcLS0tePLJJyGEwMaNG9tt6+s/39OnT5f/PGrUKIwePRr33HMPiouLMWHCBB1HdnvIycnB008/3eEGAL5+n5HxGWbmfuDAgbBYLKipqVHU19TUICQkRPWckJAQTe17qvT0dHz88cf49NNPERYWpunc3r1746GHHsKZM2e6aHTGNmDAANx3330e3z/vsR+cPXsWe/fuxZw5czSd5+v3WNu9ouU+upnPw56oLbE/e/YsCgsL2521V9PRz3dPd/fdd2PgwIEe3z/vsx/85S9/wenTpzV/vgG8z8h4DJPc+/n5YezYsSgqKpLrJElCUVGRYhbwx2JjYxXtAaCwsNBj+55GCIH09HTk5eVh3759iIyM1NyHy+XC8ePHMXjw4C4YofHV19fjq6++8vj+ff0e+7GtW7di0KBBSElJ0XSer99jkZGRCAkJUdxHTqcThw4d8ngf3cznYU/Tlth/+eWX2Lt3L+68807NfXT0893TnT9/Ht9++63H98/77AdbtmzB2LFjERUVpflcX7/PyID0/kbvj7333nvCarWK3NxccerUKfHcc8+JAQMGCIfDIYQQ4plnnhFLliyR2//1r38VvXr1En/4wx/E3//+d5GZmSl69+4tjh8/rtdb6Fbz588XgYGBori4WFRXV8vl6tWrcpsbr9nLL78s9uzZI7766itRVlYmpk+fLvz9/cXJkyf1eAvd7ne/+50oLi4WlZWV4q9//atISEgQAwcOFLW1tUII3mOeuFwuERERIRYvXux2jPeYEFeuXBHHjh0Tx44dEwDE6tWrxbFjx+SdXV577TUxYMAA8eGHH4rPP/9cTJ48WURGRorGxka5j/Hjx4v169fLrzv6PLzdtXfNmpubxRNPPCHCwsJERUWF4vOtqalJ7uPGa9bRz/ftrr1rduXKFfHSSy+J0tJSUVlZKfbu3Ssefvhhce+994pr167JffA+U/5sCiFEXV2d6Nu3r9i4caNqH752n9Htz1DJvRBCrF+/XkRERAg/Pz8RHR0tDh48KB/7xS9+IVJTUxXtd+7cKe677z7h5+cnHnzwQZGfn9/NI9YPANWydetWuc2N12zBggXy9Q0ODhaPP/64KC8v7/7B62TatGli8ODBws/PT/zkJz8R06ZNE2fOnJGP8x5Tt2fPHgFAnD592u0Y7zEhPv30U9WfxbbrIkmSWLZsmQgODhZWq1VMmDDB7VoOGTJEZGZmKura+zy83bV3zSorKz1+vn366adyHzdes45+vm937V2zq1eviokTJ4q77rpL9O7dWwwZMkTMnTvXLUnnfab82RRCiDfffFP06dNHXL58WbUPX7vP6PZnEkKILv3VABERERERdQvDrLknIiIiIqJbw+SeiIiIiKiHYHJPRERERNRDMLknIiIiIuohmNwTEREREfUQTO6JiIiIiHoIJvdERERERD0Ek3siIiIioh6CyT0RERERUQ/B5J6IiIiIqIdgck9ERERE1EP8/3gZ+SI9SdNcAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 1000x200 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "r_mix_h = numpyro.diagnostics.split_gelman_rubin(\n",
    "    mix_h.reshape((1,) + mix_h.shape))\n",
    "figure(figsize=(10,2)); imshow(r_mix_h); colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16adc9a5-0e2a-4435-809f-47651917870b",
   "metadata": {},
   "source": [
    "Seems like we have enough samples?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "id": "7e84d829-27de-4b01-9da8-8726151eea9f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "352.6536287296914"
      ]
     },
     "execution_count": 170,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "numpyro.diagnostics.effective_sample_size(\n",
    "    mix_h.reshape((1,) + mix_h.shape)).min()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28ac686f-169c-4fee-af4b-842ba7aad9ec",
   "metadata": {},
   "source": [
    "Next let's check posterior z-score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "id": "e15296f0-16e1-4169-a828-468a1a447aef",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAzYAAADFCAYAAACGoWdrAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAQqklEQVR4nO3df2xVZ/0H8M8FXGGzLTKgpVLGj7lp/MEStjUgMTSS0WZBMbpki9GyLNMsbJFUnWA2oVPDFDMJQob+4dgSdTOagXHxZyPwh4M5J+pMWFaE0MHoNkjbtcaLoff7x9fd77fj5x337twHXq/kJNxznva8c/NQ+ubc85xcoVAoBAAAQMLGZB0AAADgQik2AABA8hQbAAAgeYoNAACQPMUGAABInmIDAAAkT7EBAACSNy7rAG82MjISR44cidra2sjlclnHAQAAMlIoFOL111+PpqamGDPm7Ndkqq7YHDlyJJqbm7OOAQAAVIne3t6YPn36WcdUXbGpra2NiP8NX1dXl3EaAAAgK4ODg9Hc3FzsCGdTdcXmjY+f1dXVKTYAAMB53aJi8QAAACB5ig0AAJA8xQYAAEieYgMAACRPsQEAAJJXdauiAcAbZq56KusIRQcfvDnrCACchSs2AABA8hQbAAAgeYoNAACQPMUGAABInmIDAAAkT7EBAACSp9gAAADJU2wAAIDkKTYAAEDyFBsAACB5ig0AAJA8xQYAAEieYgMAACRPsQEAAJKn2AAAAMlTbAAAgOQpNgAAQPIUGwAAIHmKDQAAkDzFBgAASJ5iAwAAJE+xAQAAkqfYAAAAyVNsAACA5Ck2AABA8hQbAAAgeSUXm127dsXSpUujqakpcrlcbNu2bdTx5cuXRy6XG7W1tbWVKy8AAMApSi42w8PDMXfu3Ni8efMZx7S1tcXLL79c3H7yk59cUEgAAICzGVfqF7S3t0d7e/tZx9TU1ERjY+N5fb98Ph/5fL74enBwsNRIAADAJa4i99js2LEjpk6dGtdee23cddddcezYsTOOXbduXdTX1xe35ubmSkQCAAAuYmUvNm1tbfHYY49Fd3d3fOtb34qdO3dGe3t7nDx58rTjV69eHQMDA8Wtt7e33JEAAICLXMkfRTuXW2+9tfjnD37wg/GhD30o5syZEzt27IiPfvSjp4yvqamJmpqacscAAAAuIRVf7nn27NkxefLk6OnpqfSpAACAS1TFi81LL70Ux44di2nTplX6VAAAwCWq5I+iDQ0Njbr6cuDAgdi7d29MmjQpJk2aFF1dXfHJT34yGhsbY//+/XHvvffG1VdfHUuWLClrcAAAgDeUXGyeffbZaG1tLb7u7OyMiIiOjo54+OGH429/+1s8+uij0d/fH01NTXHTTTfF17/+dffRAAAAFVNysVm0aFEUCoUzHv/Nb35zQYEAAABKVfF7bAAAACpNsQEAAJKn2AAAAMlTbAAAgOQpNgAAQPIUGwAAIHmKDQAAkDzFBgAASJ5iAwAAJE+xAQAAkqfYAAAAyVNsAACA5Ck2AABA8hQbAAAgeYoNAACQPMUGAABInmIDAAAkT7EBAACSp9gAAADJU2wAAIDkKTYAAEDyFBsAACB5ig0AAJC8cVkHAADSNnPVU1lHKDr44M1ZRwAy4ooNAACQPMUGAABInmIDAAAkr+Ris2vXrli6dGk0NTVFLpeLbdu2jTpeKBTia1/7WkybNi0mTJgQixcvjhdffLFceQEAAE5RcrEZHh6OuXPnxubNm097/Nvf/nZs3LgxtmzZEnv27IkrrrgilixZEv/+978vOCwAAMDplLwqWnt7e7S3t5/2WKFQiA0bNsR9990XH//4xyMi4rHHHouGhobYtm1b3HrrrReWFgAA4DTKeo/NgQMH4ujRo7F48eLivvr6+mhpaYmnn376tF+Tz+djcHBw1AYAAFCKsj7H5ujRoxER0dDQMGp/Q0ND8dibrVu3Lrq6usoZA+CcPHcDAC4uma+Ktnr16hgYGChuvb29WUcCAAASU9Zi09jYGBERfX19o/b39fUVj71ZTU1N1NXVjdoAAABKUdZiM2vWrGhsbIzu7u7ivsHBwdizZ0/Mnz+/nKcCAAAoKvkem6Ghoejp6Sm+PnDgQOzduzcmTZoUM2bMiJUrV8Y3vvGNeM973hOzZs2K+++/P5qammLZsmXlzA0AAFBUcrF59tlno7W1tfi6s7MzIiI6Ojpi69atce+998bw8HB87nOfi/7+/li4cGH8+te/jvHjx5cvNQAAwP9TcrFZtGhRFAqFMx7P5XLxwAMPxAMPPHBBwQAAAM5X5quiAQAAXCjFBgAASJ5iAwAAJE+xAQAAkqfYAAAAyVNsAACA5Ck2AABA8hQbAAAgeYoNAACQvHFZBwCgusxc9VTWEQCgZK7YAAAAyVNsAACA5Ck2AABA8hQbAAAgeYoNAACQPMUGAABInuWeASAxluQGOJUrNgAAQPIUGwAAIHmKDQAAkDzFBgAASJ5iAwAAJE+xAQAAkqfYAAAAyfMcGwA4D54dQ6mqac4cfPDmrCNAxbliAwAAJE+xAQAAkqfYAAAAySt7sVm7dm3kcrlR23vf+95ynwYAAKCoIosHvP/974/f//73/3eScdYoAAAAKqcijWPcuHHR2NhYiW8NAABwiorcY/Piiy9GU1NTzJ49Oz796U/HoUOHzjg2n8/H4ODgqA0AAKAUuUKhUCjnN/zVr34VQ0NDce2118bLL78cXV1dcfjw4Xj++eejtrb2lPFr166Nrq6uU/YPDAxEXV1dOaMBFFXT8yUAKs1zbEjV4OBg1NfXn1c3KHuxebP+/v646qqr4qGHHoo77rjjlOP5fD7y+Xzx9eDgYDQ3Nys2QEUpNsClRLEhVaUUm4rf1T9x4sS45pproqen57THa2pqoqamptIxAACAi1jFn2MzNDQU+/fvj2nTplX6VAAAwCWq7MXmS1/6UuzcuTMOHjwYf/zjH+MTn/hEjB07Nm677bZynwoAACAiKvBRtJdeeiluu+22OHbsWEyZMiUWLlwYu3fvjilTppT7VAAAABFRgWLz+OOPl/tbAgAAnFXFFw8AAIA3VNuqlNW0Ylw1vTfV9L6cr4ovHgAAAFBpig0AAJA8xQYAAEieYgMAACRPsQEAAJKn2AAAAMlTbAAAgOR5jg2UWTWtQV9tUlwTHwBIgys2AABA8hQbAAAgeYoNAACQPMUGAABInmIDAAAkT7EBAACSp9gAAADJ8xwbLgqeHQMAZ+bfyTPz3lw8XLEBAACSp9gAAADJU2wAAIDkKTYAAEDyFBsAACB5ig0AAJA8xQYAAEie59ich2pa3/zggzdnHaGomt4X0mDOAACV4ooNAACQPMUGAABInmIDAAAkr2LFZvPmzTFz5swYP358tLS0xDPPPFOpUwEAAJe4ihSbJ554Ijo7O2PNmjXx3HPPxdy5c2PJkiXxyiuvVOJ0AADAJa4iq6I99NBDceedd8btt98eERFbtmyJp556Kn74wx/GqlWrRo3N5/ORz+eLrwcGBiIiYnBwsBLR3pKR/L+yjlDkfQEAoNKq5XfON3IUCoVzjs0VzmdUCU6cOBGXX355/OxnP4tly5YV93d0dER/f39s37591Pi1a9dGV1dXOSMAAAAXkd7e3pg+ffpZx5T9is1rr70WJ0+ejIaGhlH7GxoaYt++faeMX716dXR2dhZfj4yMxPHjx+PKK6+MXC5X7nhvyeDgYDQ3N0dvb2/U1dVlHYcEmDOUypyhVOYMpTJnKEW1zJdCoRCvv/56NDU1nXNs5g/orKmpiZqamlH7Jk6cmE2Yc6irq/ODgJKYM5TKnKFU5gylMmcoRTXMl/r6+vMaV/bFAyZPnhxjx46Nvr6+Ufv7+vqisbGx3KcDAAAof7G57LLLYt68edHd3V3cNzIyEt3d3TF//vxynw4AAKAyH0Xr7OyMjo6OuP766+PGG2+MDRs2xPDwcHGVtNTU1NTEmjVrTvnIHJyJOUOpzBlKZc5QKnOGUqQ4X8q+KtobNm3aFOvXr4+jR4/GddddFxs3boyWlpZKnAoAALjEVazYAAAAvF3Kfo8NAADA202xAQAAkqfYAAAAyVNsAACA5Ck2b8HHPvaxmDFjRowfPz6mTZsWn/nMZ+LIkSNZx6IKHTx4MO64446YNWtWTJgwIebMmRNr1qyJEydOZB2NKvbNb34zFixYEJdffnlMnDgx6zhUoc2bN8fMmTNj/Pjx0dLSEs8880zWkahiu3btiqVLl0ZTU1PkcrnYtm1b1pGoYuvWrYsbbrghamtrY+rUqbFs2bJ44YUXso51XhSbt6C1tTV++tOfxgsvvBA///nPY//+/fGpT30q61hUoX379sXIyEh8//vfj3/84x/x3e9+N7Zs2RJf/epXs45GFTtx4kTccsstcdddd2UdhSr0xBNPRGdnZ6xZsyaee+65mDt3bixZsiReeeWVrKNRpYaHh2Pu3LmxefPmrKOQgJ07d8aKFSti9+7d8bvf/S7+85//xE033RTDw8NZRzsnyz2XwS9+8YtYtmxZ5PP5eMc73pF1HKrc+vXr4+GHH45//vOfWUehym3dujVWrlwZ/f39WUehirS0tMQNN9wQmzZtioiIkZGRaG5ujnvuuSdWrVqVcTqqXS6XiyeffDKWLVuWdRQS8eqrr8bUqVNj586d8ZGPfCTrOGflis0FOn78ePzoRz+KBQsWKDWcl4GBgZg0aVLWMYAEnThxIv785z/H4sWLi/vGjBkTixcvjqeffjrDZMDFamBgICIiid9dFJu36Ctf+UpcccUVceWVV8ahQ4di+/btWUciAT09PfG9730vPv/5z2cdBUjQa6+9FidPnoyGhoZR+xsaGuLo0aMZpQIuViMjI7Fy5cr48Ic/HB/4wAeyjnNOis1/rVq1KnK53Fm3ffv2Fcd/+ctfjr/85S/x29/+NsaOHRuf/exnw6f6Lh2lzpeIiMOHD0dbW1vccsstceedd2aUnKy8lTkDAFlasWJFPP/88/H4449nHeW8jMs6QLX44he/GMuXLz/rmNmzZxf/PHny5Jg8eXJcc8018b73vS+am5tj9+7dMX/+/AonpRqUOl+OHDkSra2tsWDBgvjBD35Q4XRUo1LnDJzO5MmTY+zYsdHX1zdqf19fXzQ2NmaUCrgY3X333fHLX/4ydu3aFdOnT886znlRbP5rypQpMWXKlLf0tSMjIxERkc/nyxmJKlbKfDl8+HC0trbGvHnz4pFHHokxY1wovRRdyM8YeMNll10W8+bNi+7u7uLN3yMjI9Hd3R133313tuGAi0KhUIh77rknnnzyydixY0fMmjUr60jnTbEp0Z49e+JPf/pTLFy4MN71rnfF/v374/777485c+a4WsMpDh8+HIsWLYqrrroqvvOd78Srr75aPOZ/VzmTQ4cOxfHjx+PQoUNx8uTJ2Lt3b0REXH311fHOd74z23BkrrOzMzo6OuL666+PG2+8MTZs2BDDw8Nx++23Zx2NKjU0NBQ9PT3F1wcOHIi9e/fGpEmTYsaMGRkmoxqtWLEifvzjH8f27dujtra2eP9efX19TJgwIeN0Z2e55xL9/e9/jy984Qvx17/+NYaHh2PatGnR1tYW9913X7z73e/OOh5VZuvWrWf8ZcNfPc5k+fLl8eijj56y/w9/+EMsWrTo7Q9E1dm0aVOsX78+jh49Gtddd11s3LgxWlpaso5FldqxY0e0traesr+joyO2bt369geiquVyudPuf+SRR875keqsKTYAAEDyfNgfAABInmIDAAAkT7EBAACSp9gAAADJU2wAAIDkKTYAAEDyFBsAACB5ig0AAJA8xQYAAEieYgMAACRPsQEAAJL3P74kFj6sdk0EAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 1000x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "z_mix = (mix_h.mean(axis=0) - mix) / mix_h.std(axis=0)\n",
    "figure(figsize=(10,2)); hist(z_mix.reshape(-1), 20);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2692d9ba-caa4-4136-a847-2a62a5edd988",
   "metadata": {},
   "source": [
    "Most values w/in (-2,2) interval which is ok."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a13112f-3119-4f5a-b3ff-181f215e5644",
   "metadata": {},
   "source": [
    "Next the variability in the linear model,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "id": "22b92221-aba8-4776-9093-d3f4fb0b9ceb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAuUAAADFCAYAAADzCYzhAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAlY0lEQVR4nO3df3RV1Z338c9NIDdBklAFEiIBAlYRlFBpzYqtI2hqCCxEOzMiD6vQDDIDi3QVMy5tuixop11xpiMd6kTgUTGOnVZwFXA94oIHUwOiASQhT1GUhx8RwkN+gDMmJJhf9+znD8vVa84NOSG55+be92utveTu7HO++x72vX7Z2WcfjzHGCAAAAIBrYtzuAAAAABDtSMoBAAAAl5GUAwAAAC4jKQcAAABcRlIOAAAAuIykHAAAAHAZSTkAAADgsiGhDmhZls6dO6fExER5PJ5QhwcAAMDXGGN08eJFpaWlKSYmvOZs29ra1NHREfTncXFxio+PD2GPBkbIk/Jz584pPT091GEBAABwBbW1tRo7dqzb3fBra2tTxvjhqm/0BW2TmpqqmpqaQZ+YhzwpT0xMlCRNv+8JxQ4N/cVLPHMp5DEva01LcC128uE6V+Ie+/EYV+JK0o0vfuZabN81ca7FPvnQMNdijz7g3m+/hjUGn0UZSEMea3QlriSd+9/jXIvdeVuLa7G9B4e7Frstxb2HYCfUuff5im13531fnOBKWElSyrQG12LXHRsd8phWW5vOPvlLf54WLjo6OlTf6NOJQ+lKSuw+g9980dIN365VR0cHSblTl5esxA6N1xAXkvIhQ6yQx/THduH9+mPHeF2JG5Pg4nuOdec9S5JniHuxXb3mQ91LGoYMcefXrUOuce/vOtbr3t+1b1iXa7HdfN8x8e4l5bFeF5Ny4877jnExx3Lzsx3jYnIZrkuLhyd6NDyxe98shWd/+yLkSTkAAADgRKfxqdPmH4edxr3J1v5GUg4AAICw1iVLnUHqIwVJOQAAAMKaJSNL3WfK7eoGK5JyAAAAhLVOY4IsXyEpBwAAAEKiwxh12CTgdnWDFUk5AAAAwpr1l2JXHylIygEAABDWuoxHnab79oddNnWDFUk5AAAAwppPHvls9iS3qxusSMoBAAAQ1jpNjDpN94fEdUbOknL16RF4JSUlmjBhguLj45WVlaWDBw/2d78AAAAASVKHYoOWSOE4Kd+8ebMKCwu1Zs0aVVVVKTMzU7m5uWpsbByI/gEAACDKGeORZVNMBK0pd5yUr127VsuWLVN+fr6mTJmiDRs2aNiwYdq0adNA9A8AAABRrsPEBi2RwlFS3tHRocrKSuXk5Hx5gpgY5eTkqKKiwvaY9vZ2NTc3BxQAAACgtzoVo07F2pQ+rcQOS47eyYULF+Tz+ZSSkhJQn5KSovr6ettjiouLlZyc7C/p6el97y0AAACijs/EBC2RYsDfSVFRkZqamvyltrZ2oEMCAAAggnTZzpLHqiuCbvR0tCXiyJEjFRsbq4aGhoD6hoYGpaam2h7j9Xrl9Xr73kMAAABEtU4zRJ0268ftHig0WDmaKY+Li9OMGTNUVlbmr7MsS2VlZcrOzu73zgEAAAA+4wlaIoXjhwcVFhZqyZIl+va3v63bb79d//Zv/6bW1lbl5+cPRP8AAAAQ5aJhptxxUr5gwQKdP39eq1evVn19vaZPn66dO3d2u/kTAAAA6A+dirHd/rBTkfNIT8dJuSQVFBSooKCgv/sCAAAAdGMpRpbNqmu7usGqT0k5AAAAECqdJlZDbJevRM5MeeT88wIAAAARqcMMCVqc2Lt3r+bNm6e0tDR5PB5t3779iseUl5frtttuk9fr1Q033KDS0tKAnz/55JPyeDwBZfLkyY76JZGUAwAAIMxZxhO0ONHa2qrMzEyVlJT0qn1NTY3mzp2rWbNmqbq6WqtWrdLDDz+sXbt2BbSbOnWq6urq/GXfvn2O+iWxfAUAAABhrssMUafNrHjXX1avNDc3B9QHe05OXl6e8vLyeh13w4YNysjI0DPPPCNJuvnmm7Vv3z795je/UW5urr/dkCFDgj6zp7eYKQcAAEBY6zSxQYskpaenKzk52V+Ki4v7JW5FRYVycnIC6nJzc1VRURFQd/z4caWlpWnixIlatGiRzpw54zgWM+UAAAAIaz5JPnVfquL7y39ra2uVlJTkr++vp8nX19d32/Y7JSVFzc3N+vzzz5WQkKCsrCyVlpbqpptuUl1dnZ566indeeed+uCDD5SYmNjrWCTlAAAACGud1hDFWt3T1k7ri/UrSUlJAUl5KH11Ocy0adOUlZWl8ePHa8uWLVq6dGmvz0NSDgAAgLDW9ZWlKoH11oDGTU1NVUNDQ0BdQ0ODkpKSlJCQYHvMiBEjdOONN+rEiROOYrmWlF/6m2bFDmsPedyE3w4LeczL6u9w71GwyYfciTvuzYH9sPTkkx+MdC12xzfce98xbe6Ns/j/6nQtduqvTrkSt2rHFFfiSpL657ezfWL53Lsl6eI3fVduNEAml/y3a7HTS2tdi135P6e7EndUtXt7ULd9fHU37V2NUQ+eD3lMX2u7nK+CDp1gO6043X3FqezsbL355psBdbt371Z2dnbQY1paWnTy5En98Ic/dBSLGz0BAAAQ1q50o2dvtbS0qLq6WtXV1ZK+2PKwurraf2NmUVGRFi9e7G+/fPlynTp1So899pg+/vhjPffcc9qyZYseeeQRf5tHH31Ue/bs0SeffKL33ntPDzzwgGJjY7Vw4UJHfWP5CgAAAMJap4lVjO0TPZ39ZvrQoUOaNWuW/3VhYaEkacmSJSotLVVdXV3AzikZGRnasWOHHnnkEa1bt05jx47VCy+8ELAd4tmzZ7Vw4UJ9+umnGjVqlL73ve9p//79GjVqlKO+kZQDAAAgrPXX8pWZM2fKmODLor7+tM7Lxxw+fDjoMa+++qqjPgRDUg4AAICw1hVkpnygb/QMJZJyAAAAhLUuK0Yxlk1Sbrl343d/IykHAABAWLPkkWXz8CC7usGKpBwAAABhrdOKlcdmprzTpm6wIikHAABAWPMpRl2m+07evgja3ZukHAAAAGHNrYcHhRJJOQAAAMJaV5DlK10sXwEAAABCo8vEyGOzfMVuSctg5fid7N27V/PmzVNaWpo8Ho+2b98+AN0CAAAAvnB5+YpdiRSOk/LW1lZlZmaqpKRkIPoDAAAABOiyYoKWSOF4+UpeXp7y8vJ63b69vV3t7e3+183NzU5DAgAAIIr5jMd2+YovmmfKnSouLlZycrK/pKenD3RIAAAARBCWr/SDoqIiNTU1+Uttbe1AhwQAAEAEYflKP/B6vfJ6vQMdBgAAABHKF2T3FV8E7b7ClogAAAAIazw8CAAAAHCZz4qRx2apii+al6+0tLToxIkT/tc1NTWqrq7Wtddeq3HjxvVr5wAAAADLirFNwK1oTsoPHTqkWbNm+V8XFhZKkpYsWaLS0tJ+6xgAAAAgSUaSMfb1kcJxUj5z5kwZu6sCAAAADACfiZG40RMAAABwj8/ySFb3mzp9NnWDFUk5AAAAwpoxHhmbnVbs6gYrknIAAACENZ8VI7H7CgAAAOAey5I8NktVLMuFzgwQknIAAACEtWhYvhI5c/4AAACISJef6GlXnNi7d6/mzZuntLQ0eTwebd++/YrHlJeX67bbbpPX69UNN9xguwV4SUmJJkyYoPj4eGVlZengwYOO+iWRlAMAACDcWR4Zm2K3I0tPWltblZmZqZKSkl61r6mp0dy5czVr1ixVV1dr1apVevjhh7Vr1y5/m82bN6uwsFBr1qxRVVWVMjMzlZubq8bGRkd9Y/kKAAAAwpoxQR4e5PDROXl5ecrLy+t1+w0bNigjI0PPPPOMJOnmm2/Wvn379Jvf/Ea5ubmSpLVr12rZsmXKz8/3H7Njxw5t2rRJP/3pT3sdy7Wk/NKlOMXIG/K419f+d8hjXjbpteGuxT43N92VuFacK2ElSWnvtrkWu2aee2/85KL1rsW+e9dS12L/19JRrsQdPbHTlbiSNO2patdiv1uX4VrszsphrsU+sTretdint33Ltdif3+rO3XS3f/v/uhJXkq6Nu+Ra7P2bQv937etw7/+ZvWGsGBmbnVYu1zU3NwfUe71eeb1Xn2dWVFQoJycnoC43N1erVq2SJHV0dKiyslJFRUX+n8fExCgnJ0cVFRWOYrF8BQAAAGHNWMGLJKWnpys5OdlfiouL+yVufX29UlJSAupSUlLU3Nyszz//XBcuXJDP57NtU19f7ygWy1cAAAAQ1q60+0ptba2SkpL89f0xSx5qJOUAAAAIa8b85cZOm3pJSkpKCkjK+0tqaqoaGhoC6hoaGpSUlKSEhATFxsYqNjbWtk1qaqqjWCxfAQAAQHgznuBlAGVnZ6usrCygbvfu3crOzpYkxcXFacaMGQFtLMtSWVmZv01vkZQDAAAgvJkeigMtLS2qrq5WdXW1pC+2PKyurtaZM2ckSUVFRVq8eLG//fLly3Xq1Ck99thj+vjjj/Xcc89py5YteuSRR/xtCgsL9fzzz+vll1/WRx99pBUrVqi1tdW/G0tvsXwFAAAA4S3YnuQO9yk/dOiQZs2a5X9dWFgoSVqyZIlKS0tVV1fnT9AlKSMjQzt27NAjjzyidevWaezYsXrhhRf82yFK0oIFC3T+/HmtXr1a9fX1mj59unbu3Nnt5s8rISkHAABAWPvqTitfr3di5syZMj1sbm73tM6ZM2fq8OHDPZ63oKBABQUFzjrzNSTlAAAACG/B1o8P8JryUCIpBwAAQFjzWF8Uu/pIQVIOAACA8NZPa8rDGUk5AAAAwluwnVYc7r4SzhxtiVhcXKzvfOc7SkxM1OjRo3X//ffr2LFjA9U3AAAA4MuZcrsSIRwl5Xv27NHKlSu1f/9+7d69W52dnbr33nvV2to6UP0DAABAlLu8ptyuRApHy1d27twZ8Lq0tFSjR49WZWWl/uqv/sr2mPb2drW3t/tfNzc396GbAAAAQOS6qid6NjU1SZKuvfbaoG2Ki4uVnJzsL+np6VcTEgAAAFHGYzzyWDYlgrZE7HNSblmWVq1ape9+97u65ZZbgrYrKipSU1OTv9TW1vY1JAAAAKKR1UOJEH3efWXlypX64IMPtG/fvh7beb1eeb3evoYBAABAlPOYL4pdfaToU1JeUFCgN954Q3v37tXYsWP7u08AAADAl4LNikfrTLkxRj/+8Y+1bds2lZeXKyMjY6D6BQAAAEiSfw25XX2kcJSUr1y5Ur///e/1+uuvKzExUfX19ZKk5ORkJSQkDEgHAQAAEOV4eFCg9evXq6mpSTNnztSYMWP8ZfPmzQPVPwAAAEQ59in/GmMi6J8jAAAAGByCJeDRmpQDAAAAIRcFy1dIygEAABDW2BIRAAAAcBsz5QAAAIC7PMZ+TTkz5QAAAECIBNtpJWp3XwEAAABCjuUrAAAAgLuYKR9Aww4NU6w3PuRx674/LOQxL0vbdsq12MPS3XnianuSo+dT9au264a6FvvaI+499vfe15a4Fjs21udabLe0pLo3t/G//pzpWmy1uffZzvg/7a7FnvTgSddiH395imuxP7vBne/TIw2TXYkrSZfGd7oW+xuuRQ5fJOUAAACA26Jg+Yp7Ux0AAABAL1yeKbcrfVFSUqIJEyYoPj5eWVlZOnjwYNC2nZ2d+sUvfqFJkyYpPj5emZmZ2rlzZ0CbJ598Uh6PJ6BMnuzsNz0k5QAAAAhr/ZmUb968WYWFhVqzZo2qqqqUmZmp3NxcNTY22rZ/4okntHHjRj377LM6evSoli9frgceeECHDx8OaDd16lTV1dX5y759+xz1i6QcAAAA4c30UBxau3atli1bpvz8fE2ZMkUbNmzQsGHDtGnTJtv2r7zyin72s59pzpw5mjhxolasWKE5c+bomWeeCWg3ZMgQpaam+svIkSMd9YukHAAAAGHNY4IXSWpubg4o7e32N4R3dHSosrJSOTk5/rqYmBjl5OSooqLC9pj29nbFxwduTpKQkNBtJvz48eNKS0vTxIkTtWjRIp05c8bReyQpBwAAQFi70vKV9PR0JScn+0txcbHteS5cuCCfz6eUlJSA+pSUFNXX19sek5ubq7Vr1+r48eOyLEu7d+/W1q1bVVdX52+TlZWl0tJS7dy5U+vXr1dNTY3uvPNOXbx4sdfvkd1XAAAAEN6usPtKbW2tkpKS/NVer7ffQq9bt07Lli3T5MmT5fF4NGnSJOXn5wcsd8nLy/P/edq0acrKytL48eO1ZcsWLV26tFdxmCkHAABAWPOYIDPlf0nKk5KSAkqwpHzkyJGKjY1VQ0NDQH1DQ4NSU1Ntjxk1apS2b9+u1tZWnT59Wh9//LGGDx+uiRMnBu3viBEjdOONN+rEiRO9fo8k5QAAAAhr/bX7SlxcnGbMmKGysjJ/nWVZKisrU3Z2do/HxsfH6/rrr1dXV5f++Mc/av78+UHbtrS06OTJkxozZkyv+0ZSDgAAgPDWj7uvFBYW6vnnn9fLL7+sjz76SCtWrFBra6vy8/MlSYsXL1ZRUZG//YEDB7R161adOnVK77zzjmbPni3LsvTYY4/52zz66KPas2ePPvnkE7333nt64IEHFBsbq4ULF/a6X6wpBwAAQFgLNivel33KFyxYoPPnz2v16tWqr6/X9OnTtXPnTv/Nn2fOnFFMzJfz1m1tbXriiSd06tQpDR8+XHPmzNErr7yiESNG+NucPXtWCxcu1KeffqpRo0bpe9/7nvbv369Ro0b1ul+OkvL169dr/fr1+uSTTyR9sUn66tWrAxa3AwAAAP2pP5NySSooKFBBQYHtz8rLywNe33XXXTp69GiP53v11Vf71pGvcLR8ZezYsXr66adVWVmpQ4cO6e6779b8+fP14YcfXnVHAAAAAFv9uHwlXDmaKZ83b17A61/96ldav3699u/fr6lTp/ZrxwAAAABJ8lhGHqt7Bm5XN1j1eU25z+fTa6+9ptbW1h7vVm1vbw94qlJzc3NfQwIAACAK9ffylXDkePeVI0eOaPjw4fJ6vVq+fLm2bdumKVOmBG1fXFwc8ISl9PT0q+owAAAAoovHBC+RwnFSftNNN6m6uloHDhzQihUrtGTJkh4XvxcVFampqclfamtrr6rDAAAAiC79tU95OHO8fCUuLk433HCDJGnGjBl6//33tW7dOm3cuNG2vdfr7ddHnQIAACDKmCAJeATNlF/1PuWWZQWsGQcAAAD6lTFfFLv6COEoKS8qKlJeXp7GjRunixcv6ve//73Ky8u1a9eugeofAAAAolw03OjpKClvbGzU4sWLVVdXp+TkZE2bNk27du3S97///YHqHwAAAKKcxyd5bO6E9PhC35eB4igpf/HFFweqHwAAAICtYDutRNLuK1e9phwAAAAYSDw8CAAAAHAZa8oBAAAAt7H7CgAAAOAuZsoBAAAAl3l8Rp4YmzXlPmbKAQAAgNAwsn96Z+Tk5CTlAAAACG8eE2T3FdaUAwAAAKHBmvIBNCynUbHXeEMe9/M3U0Ie0y9uqGuhP7sh1pW47de69y/YjNfbXYvdmjrMtdhDz5x3LXbX/zvnWuzTRXe4EjfWvWGmubcecS32wusqXIv9y+f+h2ux330j07XYKVaHa7GfWbXRlbgP78l3Ja4kDWl07//Zf7PiTyGP2dbSqQ9eCHnYXvMYYzsrzkw5AAAAECIen5HH5vGd3OgJAAAAhIplvih29RGCpBwAAABhzWO+KHb1kYKkHAAAAGGN5SsAAACA26Jg+UqM2x0AAAAAenJ59xW70hclJSWaMGGC4uPjlZWVpYMHDwZt29nZqV/84heaNGmS4uPjlZmZqZ07d17VOe2QlAMAACC8WUby2ZQ+zJRv3rxZhYWFWrNmjaqqqpSZmanc3Fw1Njbatn/iiSe0ceNGPfvsszp69KiWL1+uBx54QIcPH+7zOe2QlAMAACCseSwjj2XZFOdJ+dq1a7Vs2TLl5+drypQp2rBhg4YNG6ZNmzbZtn/llVf0s5/9THPmzNHEiRO1YsUKzZkzR88880yfz2mHpBwAAADhzZjgRVJzc3NAaW+3f7JbR0eHKisrlZOT46+LiYlRTk6OKirsH4rW3t6u+Pj4gLqEhATt27evz+e0Q1IOAACAsObxmaBFktLT05WcnOwvxcXFtue5cOGCfD6fUlICn/CekpKi+vp622Nyc3O1du1aHT9+XJZlaffu3dq6davq6ur6fE477L4CAACA8GZZkseyr5dUW1urpKQkf7XX6+230OvWrdOyZcs0efJkeTweTZo0Sfn5+Y6WpvTGVc2UP/300/J4PFq1alU/dQcAAAD4missX0lKSgoowZLykSNHKjY2Vg0NDQH1DQ0NSk1NtT1m1KhR2r59u1pbW3X69Gl9/PHHGj58uCZOnNjnc9rpc1L+/vvva+PGjZo2bVpfTwEAAABc0ZWWr/RWXFycZsyYobKyMn+dZVkqKytTdnZ2j8fGx8fr+uuvV1dXl/74xz9q/vz5V33Or+pTUt7S0qJFixbp+eef1ze+8Y2+nAIAAADoHZ8VvDhUWFio559/Xi+//LI++ugjrVixQq2trcrPz5ckLV68WEVFRf72Bw4c0NatW3Xq1Cm98847mj17tizL0mOPPdbrc/ZGn9aUr1y5UnPnzlVOTo5++ctf9ti2vb094A7Y5ubmvoQEAABAtPrKUpVu9Q4tWLBA58+f1+rVq1VfX6/p06dr586d/hs1z5w5o5iYL+et29ra9MQTT+jUqVMaPny45syZo1deeUUjRozo9Tl7w3FS/uqrr6qqqkrvv/9+r9oXFxfrqaeechoGAAAA+IKx/Dd1dqvvg4KCAhUUFNj+rLy8POD1XXfdpaNHj17VOXvD0fKV2tpa/eQnP9F//ud/dtuvMZiioiI1NTX5S21tbZ86CgAAgCjl8wUvEcLRTHllZaUaGxt12223+et8Pp/27t2rf//3f1d7e7tiY2MDjvF6vf26LQ0AAACiTD8uXwlXjpLye+65R0eOHAmoy8/P1+TJk/X44493S8gBAACAq+az7Jeq2C1pGaQcJeWJiYm65ZZbAuquueYaXXfddd3qAQAAgH5hGUl2SXmUzpQDAAAAIWdZsk/Ko3Sm3M7X71AFAAAA+hVJOQAAAOAu4/PJmO47rRgrSndfAQAAAELOGPv149G6+woAAAAQcj6f5LGZFbeZPR+sSMoBAAAQ1ozPJ2OTlNstaRmsSMoBAAAQ3oyRxPIVAAAAwD0+K8jyFXZfAQAAAELii+UrMd3rWb7Sd+Yvv2bwXWoPdegv4ra3uRJXkrosd96z5N77ttrc+7VSV5d7f9e+ju5fHKHi5jjrMp2uxXbts+3e5VZHi3vXuzXOvdmpLl/0fZdKUldXh2uxWy+6k/hYn7t3va0295K9Nhc+2+2tXZK+zNPCTafVIWOzfKVL7n0P9jePCfHVP3v2rNLT00MZEgAAAL1QW1ursWPHut0Nv7a2NmVkZKi+vj5om9TUVNXU1Cg+Pj6EPet/IU/KLcvSuXPnlJiYKI/H4+jY5uZmpaenq7a2VklJSQPUw8jCNXOOa+Yc18w5rplzXDPnuGbORes1M8bo4sWLSktLU0yMe7/ttdPW1qaOjuC/KYqLixv0CbnkwvKVmJiYq/4XWFJSUlR9UPoD18w5rplzXDPnuGbOcc2c45o5F43XLDk52e0u2IqPj4+IpPtKwuufQgAAAEAUIikHAAAAXDaoknKv16s1a9bI6/W63ZVBg2vmHNfMOa6Zc1wz57hmznHNnOOawS0hv9ETAAAAQKBBNVMOAAAARCKScgAAAMBlJOUAAACAy0jKAQAAAJeRlAMAAAAuC7ukvKSkRBMmTFB8fLyysrJ08ODBHtu/9tprmjx5suLj43XrrbfqzTffDFFP3VdcXKzvfOc7SkxM1OjRo3X//ffr2LFjPR5TWloqj8cTUKLhKVmXPfnkk93e/+TJk3s8JprHmCRNmDCh2zXzeDxauXKlbftoHGN79+7VvHnzlJaWJo/Ho+3btwf83Bij1atXa8yYMUpISFBOTo6OHz9+xfM6/T4cTHq6Zp2dnXr88cd166236pprrlFaWpoWL16sc+fO9XjOvny+B5MrjbMf/ehH3d7/7Nmzr3jeaB1nkmy/2zwej379618HPWekjzO4J6yS8s2bN6uwsFBr1qxRVVWVMjMzlZubq8bGRtv27733nhYuXKilS5fq8OHDuv/++3X//ffrgw8+CHHP3bFnzx6tXLlS+/fv1+7du9XZ2al7771Xra2tPR6XlJSkuro6fzl9+nSIehwepk6dGvD+9+3bF7RttI8xSXr//fcDrtfu3bslSX/7t38b9JhoG2Otra3KzMxUSUmJ7c//5V/+Rb/97W+1YcMGHThwQNdcc41yc3PV1tYW9JxOvw8Hm56u2aVLl1RVVaWf//znqqqq0tatW3Xs2DHdd999Vzyvk8/3YHOlcSZJs2fPDnj/f/jDH3o8ZzSPM0kB16qurk6bNm2Sx+PRX//1X/d43kgeZ3CRCSO33367Wblypf+1z+czaWlppri42Lb9gw8+aObOnRtQl5WVZf7hH/5hQPsZrhobG40ks2fPnqBtXnrpJZOcnBy6ToWZNWvWmMzMzF63Z4x195Of/MRMmjTJWJZl+/NoH2OSzLZt2/yvLcsyqamp5te//rW/7rPPPjNer9f84Q9/CHoep9+Hg9nXr5mdgwcPGknm9OnTQds4/XwPZnbXbMmSJWb+/PmOzsM4CzR//nxz991399gmmsYZQitsZso7OjpUWVmpnJwcf11MTIxycnJUUVFhe0xFRUVAe0nKzc0N2j7SNTU1SZKuvfbaHtu1tLRo/PjxSk9P1/z58/Xhhx+Gonth4/jx40pLS9PEiRO1aNEinTlzJmhbxligjo4O/e53v9Pf/d3fyePxBG0X7WPsq2pqalRfXx8wjpKTk5WVlRV0HPXl+zDSNTU1yePxaMSIET22c/L5jkTl5eUaPXq0brrpJq1YsUKffvpp0LaMs0ANDQ3asWOHli5desW20T7OMDDCJim/cOGCfD6fUlJSAupTUlJUX19ve0x9fb2j9pHMsiytWrVK3/3ud3XLLbcEbXfTTTdp06ZNev311/W73/1OlmXpjjvu0NmzZ0PYW/dkZWWptLRUO3fu1Pr161VTU6M777xTFy9etG3PGAu0fft2ffbZZ/rRj34UtE20j7GvuzxWnIyjvnwfRrK2tjY9/vjjWrhwoZKSkoK2c/r5jjSzZ8/Wf/zHf6isrEz//M//rD179igvL08+n8+2PeMs0Msvv6zExET94Ac/6LFdtI8zDJwhbncA/WPlypX64IMPrriuLTs7W9nZ2f7Xd9xxh26++WZt3LhR//RP/zTQ3XRdXl6e/8/Tpk1TVlaWxo8fry1btvRqdiTavfjii8rLy1NaWlrQNtE+xtC/Ojs79eCDD8oYo/Xr1/fYNto/3w899JD/z7feequmTZumSZMmqby8XPfcc4+LPRscNm3apEWLFl3xxvRoH2cYOGEzUz5y5EjFxsaqoaEhoL6hoUGpqam2x6SmpjpqH6kKCgr0xhtv6O2339bYsWMdHTt06FB961vf0okTJwaod+FtxIgRuvHGG4O+f8bYl06fPq233npLDz/8sKPjon2MXR4rTsZRX74PI9HlhPz06dPavXt3j7Pkdq70+Y50EydO1MiRI4O+f8bZl9555x0dO3bM8febxDhD/wmbpDwuLk4zZsxQWVmZv86yLJWVlQXMun1VdnZ2QHtJ2r17d9D2kcYYo4KCAm3btk1/+tOflJGR4fgcPp9PR44c0ZgxYwagh+GvpaVFJ0+eDPr+o32MfdVLL72k0aNHa+7cuY6Oi/YxlpGRodTU1IBx1NzcrAMHDgQdR335Pow0lxPy48eP66233tJ1113n+BxX+nxHurNnz+rTTz8N+v4ZZ1968cUXNWPGDGVmZjo+NtrHGfqR23eaftWrr75qvF6vKS0tNUePHjV///d/b0aMGGHq6+uNMcb88Ic/ND/96U/97d99910zZMgQ86//+q/mo48+MmvWrDFDhw41R44ccesthNSKFStMcnKyKS8vN3V1df5y6dIlf5uvX7OnnnrK7Nq1y5w8edJUVlaahx56yMTHx5sPP/zQjbcQcv/4j/9oysvLTU1NjXn33XdNTk6OGTlypGlsbDTGMMaC8fl8Zty4cebxxx/v9jPGmDEXL140hw8fNocPHzaSzNq1a83hw4f9O4U8/fTTZsSIEeb11183f/7zn838+fNNRkaG+fzzz/3nuPvuu82zzz7rf32l78PBrqdr1tHRYe677z4zduxYU11dHfD91t7e7j/H16/ZlT7fg11P1+zixYvm0UcfNRUVFaampsa89dZb5rbbbjPf/OY3TVtbm/8cjLPAz6YxxjQ1NZlhw4aZ9evX254j2sYZ3BNWSbkxxjz77LNm3LhxJi4uztx+++1m//79/p/dddddZsmSJQHtt2zZYm688UYTFxdnpk6danbs2BHiHrtHkm156aWX/G2+fs1WrVrlv74pKSlmzpw5pqqqKvSdd8mCBQvMmDFjTFxcnLn++uvNggULzIkTJ/w/Z4zZ27Vrl5Fkjh071u1njDFj3n77bdvP4uXrYlmW+fnPf25SUlKM1+s199xzT7drOX78eLNmzZqAup6+Dwe7nq5ZTU1N0O+3t99+23+Or1+zK32+B7uertmlS5fMvffea0aNGmWGDh1qxo8fb5YtW9YtuWacBX42jTFm48aNJiEhwXz22We254i2cQb3eIwxZkCn4gEAAAD0KGzWlAMAAADRiqQcAAAAcBlJOQAAAOAyknIAAADAZSTlAAAAgMtIygEAAACXkZQDAAAALiMpBwAAAFxGUg4AAAC4jKQcAAAAcBlJOQAAAOCy/w+2GJKzxtWVjAAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 1000x200 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=(10,2)); imshow(mix_h.std(axis=0));  colorbar();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2d44047-8ad4-486f-b5ad-07ecb6cd6902",
   "metadata": {},
   "source": [
    "First note that the prior standard deviation on this was 1, so there's no significant shrinkage."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "802a8a35-7f32-4980-8982-d2867a2df4b0",
   "metadata": {},
   "source": [
    "## variational"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "id": "e40c60e1-0270-43e7-b54e-4109a8325453",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "guide = numpyro.infer.autoguide.AutoDiagonalNormal(log_p)\n",
    "optimizer = numpyro.optim.Adam(step_size=1e-3)\n",
    "svi = numpyro.infer.SVI(log_p, guide, optimizer, \n",
    "                        loss=numpyro.infer.Trace_ELBO())\n",
    "svi_result = svi.run(next_key(), 2000, progress_bar=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e30534a0-69be-4edf-93dc-f370fb5d607f",
   "metadata": {},
   "source": [
    "With mean field VI we have worse z scores,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "id": "a776d548-daad-4237-9910-cd397a179f45",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAzYAAADFCAYAAACGoWdrAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAVWklEQVR4nO3dfWxV9f0H8E8RaGVSOqC2dPLoAzgV5nA2ZT7MSEQkmzjjlBFlhvk0dGM4N1g2ETd/MDVqdGa6ZMKWLVFcNt2G0yAKRkFUhEwQjRAQEApRQys4QeD7+2PxboW2WL23l2Nfr+Qm3HM+557Pvd9+6X333HNuSUopBQAAQIZ1KnYDAAAAn5ZgAwAAZJ5gAwAAZJ5gAwAAZJ5gAwAAZJ5gAwAAZJ5gAwAAZF7nYjewv3379sXmzZuje/fuUVJSUux2AACAIkkpxXvvvRc1NTXRqVPrx2QOuWCzefPm6Nu3b7HbAAAADhEbN26Mo446qtWaQy7YdO/ePSL+03x5eXmRuwEAAIqlsbEx+vbtm8sIrTnkgs1HHz8rLy8XbAAAgI91ioqLBwAAAJkn2AAAAJkn2AAAAJkn2AAAAJkn2AAAAJl3yF0VDQDIlgFT5xW7hZz1s8YUuwWgSByxAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMk+wAQAAMq9zsRsAgCwYMHVesVvIWT9rTLFbADjkOGIDAABknmADAABknmADAABknmADAABknmADAABknmADAABknmADAABknmADAABknmADAABkXudiNwAAtM2AqfOK3QLAIccRGwAAIPMEGwAAIPMEGwAAIPMEGwAAIPMEGwAAIPMEGwAAIPPaHGyeeeaZ+PrXvx41NTVRUlISjzzySJP1KaW48cYbo0+fPnH44YfHyJEj44033shXvwAAAAdoc7DZuXNnDBs2LO69995m1996661x9913x3333RdLly6Nz33uczFq1Kj44IMPPnWzAAAAzWnzF3SOHj06Ro8e3ey6lFLcdddd8bOf/SzOP//8iIj4wx/+EFVVVfHII4/EJZdc8um6BQAAaEZez7FZt25d1NfXx8iRI3PLevToEbW1tbFkyZJmt9m1a1c0NjY2uQEAALRFXoNNfX19RERUVVU1WV5VVZVbt7+ZM2dGjx49cre+ffvmsyUAAKADKPpV0aZNmxYNDQ2528aNG4vdEgAAkDF5DTbV1dUREbF169Ymy7du3Zpbt7/S0tIoLy9vcgMAAGiLvAabgQMHRnV1dSxYsCC3rLGxMZYuXRp1dXX53BUAAEBOm6+KtmPHjlizZk3u/rp162LFihXRs2fP6NevX0yePDl++ctfxrHHHhsDBw6Mn//851FTUxNjx47NZ98AAAA5bQ42L730Upx11lm5+1OmTImIiAkTJsScOXPixz/+cezcuTOuvPLK2L59e5x22mnx+OOPR1lZWf66BgAA+B8lKaVU7Cb+V2NjY/To0SMaGhqcbwPAIWPA1HnFboGPYf2sMcVuAcijtmSDol8VDQAA4NMSbAAAgMwTbAAAgMxr88UDAPhsO5TOJXG+BAAflyM2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5gk2AABA5nUudgMA0JIBU+cVuwUyxs9M89bPGlPsFqDgHLEBAAAyT7ABAAAyT7ABAAAyT7ABAAAyT7ABAAAyT7ABAAAyT7ABAAAyT7ABAAAyT7ABAAAyr3OxGwDo6HxTOgB8eo7YAAAAmSfYAAAAmSfYAAAAmSfYAAAAmSfYAAAAmSfYAAAAmZf3YHPTTTdFSUlJk9uQIUPyvRsAAICcgnyPzQknnBBPPvnkf3fS2dflAAAAhVOQxNG5c+eorq4uxEMDAAAcoCDn2LzxxhtRU1MTgwYNivHjx8eGDRtarN21a1c0NjY2uQEAALRF3o/Y1NbWxpw5c2Lw4MGxZcuWmDFjRpx++umxcuXK6N69+wH1M2fOjBkzZuS7DQAAOKgBU+cVu4Wc9bPGFLuFTMv7EZvRo0fHRRddFEOHDo1Ro0bFY489Ftu3b4+5c+c2Wz9t2rRoaGjI3TZu3JjvlgAAgM+4gp/VX1FREccdd1ysWbOm2fWlpaVRWlpa6DYAAIDPsIJ/j82OHTti7dq10adPn0LvCgAA6KDyHmx+9KMfxaJFi2L9+vWxePHiuOCCC+Kwww6LcePG5XtXAAAAEVGAj6Jt2rQpxo0bF++8805UVlbGaaedFs8//3xUVlbme1cAAAARUYBg8+CDD+b7IQEAAFpV8HNsAAAACk2wAQAAMk+wAQAAMq/g32MDAEBxDZg6r9gtQME5YgMAAGSeYAMAAGSeYAMAAGSeYAMAAGSeYAMAAGSeYAMAAGSeYAMAAGSeYAMAAGSeYAMAAGSeYAMAAGRe52I3ABTWgKnzit1CzvpZY4rdQs6h9LoAAJ+eIzYAAEDmCTYAAEDmCTYAAEDmCTYAAEDmCTYAAEDmCTYAAEDmCTYAAEDmCTYAAEDmCTYAAEDmdS52A0DHMWDqvGK3AACHrEPp9+T6WWOK3UKbOWIDAABknmADAABknmADAABknmADAABknmADAABknmADAABknmADAABknmADAABknmADAABknmADAABkXudiN5AFA6bOK3YLHMT6WWOK3UKOnxcAgPbniA0AAJB5gg0AAJB5gg0AAJB5gg0AAJB5gg0AAJB5gg0AAJB5BQs29957bwwYMCDKysqitrY2XnjhhULtCgAA6OAKEmweeuihmDJlSkyfPj1efvnlGDZsWIwaNSq2bdtWiN0BAAAdXEG+oPOOO+6IK664Ii6//PKIiLjvvvti3rx58cADD8TUqVOb1O7atSt27dqVu9/Q0BAREY2NjYVo7RPZt+v9YrfAQfh5AQDIn0PlvdVHfaSUDlpbkj5OVRvs3r07unXrFn/+859j7NixueUTJkyI7du3x6OPPtqk/qabbooZM2bkswUAAOAzZOPGjXHUUUe1WpP3IzZvv/127N27N6qqqposr6qqitdee+2A+mnTpsWUKVNy9/ft2xfvvvtu9OrVK0pKSvLdXlE0NjZG3759Y+PGjVFeXl7sdmiGMcoG45QNxunQZ4yywThlg3EqrJRSvPfee1FTU3PQ2oJ8FK0tSktLo7S0tMmyioqK4jRTYOXl5X7gD3HGKBuMUzYYp0OfMcoG45QNxqlwevTo8bHq8n7xgN69e8dhhx0WW7dubbJ869atUV1dne/dAQAA5D/YdO3aNYYPHx4LFizILdu3b18sWLAg6urq8r07AACAwnwUbcqUKTFhwoQ45ZRT4tRTT4277rordu7cmbtKWkdTWloa06dPP+Ajdxw6jFE2GKdsME6HPmOUDcYpG4zToSPvV0X7yK9//eu47bbbor6+Pr70pS/F3XffHbW1tYXYFQAA0MEVLNgAAAC0l7yfYwMAANDeBBsAACDzBBsAACDzBBsAACDzBJs8uOWWW2LEiBHRrVu3qKioaLZmw4YNMWbMmOjWrVsceeSRccMNN8SePXtafdx33303xo8fH+Xl5VFRURETJ06MHTt2FOAZdDwLFy6MkpKSZm8vvvhii9t97WtfO6D+6quvbsfOO54BAwYc8JrPmjWr1W0++OCDmDRpUvTq1SuOOOKIuPDCCw/40mDyY/369TFx4sQYOHBgHH744XH00UfH9OnTY/fu3a1uZy4V3r333hsDBgyIsrKyqK2tjRdeeKHV+ocffjiGDBkSZWVlcdJJJ8Vjjz3WTp12TDNnzoyvfOUr0b179zjyyCNj7Nix8frrr7e6zZw5cw6YN2VlZe3Uccd00003HfCaDxkypNVtzKXiEWzyYPfu3XHRRRfFNddc0+z6vXv3xpgxY2L37t2xePHi+P3vfx9z5syJG2+8sdXHHT9+fKxatSrmz58f//jHP+KZZ56JK6+8shBPocMZMWJEbNmypcntu9/9bgwcODBOOeWUVre94oormmx36623tlPXHdfNN9/c5DW/7rrrWq3/4Q9/GH//+9/j4YcfjkWLFsXmzZvjm9/8Zjt127G89tprsW/fvrj//vtj1apVceedd8Z9990XP/3pTw+6rblUOA899FBMmTIlpk+fHi+//HIMGzYsRo0aFdu2bWu2fvHixTFu3LiYOHFiLF++PMaOHRtjx46NlStXtnPnHceiRYti0qRJ8fzzz8f8+fPjww8/jHPOOSd27tzZ6nbl5eVN5s2bb77ZTh13XCeccEKT1/zZZ59tsdZcKrJE3syePTv16NHjgOWPPfZY6tSpU6qvr88t+81vfpPKy8vTrl27mn2sV199NUVEevHFF3PL/vnPf6aSkpL01ltv5b33jm737t2psrIy3Xzzza3WnXnmmekHP/hB+zRFSiml/v37pzvvvPNj12/fvj116dIlPfzww7llq1evThGRlixZUoAO2d+tt96aBg4c2GqNuVRYp556apo0aVLu/t69e1NNTU2aOXNms/Xf+ta30pgxY5osq62tTVdddVVB++S/tm3bliIiLVq0qMWalt5nUDjTp09Pw4YN+9j15lJxOWLTDpYsWRInnXRSVFVV5ZaNGjUqGhsbY9WqVS1uU1FR0eTowciRI6NTp06xdOnSgvfc0fztb3+Ld955Jy6//PKD1v7pT3+K3r17x4knnhjTpk2L999/vx067NhmzZoVvXr1ipNPPjluu+22Vj/GuWzZsvjwww9j5MiRuWVDhgyJfv36xZIlS9qj3Q6voaEhevbsedA6c6kwdu/eHcuWLWsyBzp16hQjR45scQ4sWbKkSX3Ef35PmTPtp6GhISLioHNnx44d0b9//+jbt2+cf/75Lb6PIH/eeOONqKmpiUGDBsX48eNjw4YNLdaaS8XVudgNdAT19fVNQk1E5O7X19e3uM2RRx7ZZFnnzp2jZ8+eLW7DJ/e73/0uRo0aFUcddVSrdd/+9rejf//+UVNTE//617/iJz/5Sbz++uvxl7/8pZ067Xi+//3vx5e//OXo2bNnLF68OKZNmxZbtmyJO+64o9n6+vr66Nq16wHnu1VVVZk77WDNmjVxzz33xO23395qnblUOG+//Xbs3bu32d87r732WrPbtPR7ypxpH/v27YvJkyfHV7/61TjxxBNbrBs8eHA88MADMXTo0GhoaIjbb789RowYEatWrTro7y8+mdra2pgzZ04MHjw4tmzZEjNmzIjTTz89Vq5cGd27dz+g3lwqLsGmBVOnTo1f/epXrdasXr36oCeQ0b4+ybht2rQpnnjiiZg7d+5BH/9/z3E66aSTok+fPnH22WfH2rVr4+ijj/7kjXcwbRmnKVOm5JYNHTo0unbtGldddVXMnDkzSktLC91qh/VJ5tJbb70V5557blx00UVxxRVXtLqtuQT/NWnSpFi5cmWr525ERNTV1UVdXV3u/ogRI+L444+P+++/P37xi18Uus0OafTo0bl/Dx06NGpra6N///4xd+7cmDhxYhE7ozmCTQuuv/76+M53vtNqzaBBgz7WY1VXVx9wNZqPrtBUXV3d4jb7n+S5Z8+eePfdd1vchk82brNnz45evXrFN77xjTbvr7a2NiL+81dqb8Y+vk8zv2pra2PPnj2xfv36GDx48AHrq6urY/fu3bF9+/YmR222bt1q7rRBW8do8+bNcdZZZ8WIESPit7/9bZv3Zy7lT+/eveOwww474EqArc2B6urqNtWTP9dee23uAkFtPerSpUuXOPnkk2PNmjUF6o79VVRUxHHHHdfia24uFZdg04LKysqorKzMy2PV1dXFLbfcEtu2bct9vGz+/PlRXl4eX/ziF1vcZvv27bFs2bIYPnx4REQ89dRTsW/fvtwbAA7U1nFLKcXs2bPjsssuiy5durR5fytWrIiIiD59+rR5247s08yvFStWRKdOnQ74qOZHhg8fHl26dIkFCxbEhRdeGBERr7/+emzYsKHJXzppXVvG6K233oqzzjorhg8fHrNnz45Ondp++qa5lD9du3aN4cOHx4IFC2Ls2LER8Z+POi1YsCCuvfbaZrepq6uLBQsWxOTJk3PL5s+fb84UUEoprrvuuvjrX/8aCxcujIEDB7b5Mfbu3RuvvPJKnHfeeQXokObs2LEj1q5dG5deemmz682lIiv21Qs+C9588820fPnyNGPGjHTEEUek5cuXp+XLl6f33nsvpZTSnj170oknnpjOOeectGLFivT444+nysrKNG3atNxjLF26NA0ePDht2rQpt+zcc89NJ598clq6dGl69tln07HHHpvGjRvX7s/vs+zJJ59MEZFWr159wLpNmzalwYMHp6VLl6aUUlqzZk26+eab00svvZTWrVuXHn300TRo0KB0xhlntHfbHcbixYvTnXfemVasWJHWrl2b/vjHP6bKysp02WWX5Wr2H6eUUrr66qtTv3790lNPPZVeeumlVFdXl+rq6orxFD7zNm3alI455ph09tlnp02bNqUtW7bkbv9bYy61rwcffDCVlpamOXPmpFdffTVdeeWVqaKiInd1zksvvTRNnTo1V//cc8+lzp07p9tvvz2tXr06TZ8+PXXp0iW98sorxXoKn3nXXHNN6tGjR1q4cGGTefP+++/navYfpxkzZqQnnngirV27Ni1btixdcsklqaysLK1ataoYT6FDuP7669PChQvTunXr0nPPPZdGjhyZevfunbZt25ZSMpcONYJNHkyYMCFFxAG3p59+Olezfv36NHr06HT44Yen3r17p+uvvz59+OGHufVPP/10ioi0bt263LJ33nknjRs3Lh1xxBGpvLw8XX755bmwRH6MGzcujRgxotl169atazKOGzZsSGeccUbq2bNnKi0tTcccc0y64YYbUkNDQzt23LEsW7Ys1dbWph49eqSysrJ0/PHHp//7v/9LH3zwQa5m/3FKKaV///vf6Xvf+176/Oc/n7p165YuuOCCJm+0yZ/Zs2c3+//f//7dzFwqjnvuuSf169cvde3aNZ166qnp+eefz60788wz04QJE5rUz507Nx133HGpa9eu6YQTTkjz5s1r5447lpbmzezZs3M1+4/T5MmTc2NaVVWVzjvvvPTyyy+3f/MdyMUXX5z69OmTunbtmr7whS+kiy++OK1Zsya33lw6tJSklFJ7HiECAADIN99jAwAAZJ5gAwAAZJ5gAwAAZJ5gAwAAZJ5gAwAAZJ5gAwAAZJ5gAwAAZJ5gAwAAZJ5gAwAAZJ5gAwAAZJ5gAwAAZN7/A3yVauO7jWXxAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 1000x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "loc = svi_result.params['auto_loc']\n",
    "scale = svi_result.params['auto_scale']\n",
    "z_mix = (loc[:100].reshape(mix.shape) - mix) / scale[:100].reshape(mix.shape)\n",
    "figure(figsize=(10,2)); hist(z_mix.reshape(-1), 20);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46f5dd98-23e5-4e35-9490-1da77edde36e",
   "metadata": {},
   "source": [
    "but better posterior shrinkage,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "id": "8d6f707f-8043-44f8-90fe-e393a65e7eff",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAtwAAADFCAYAAABn7yUTAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAfNklEQVR4nO3df3DU1f3v8dduIBt+JKGYkhCJ/LKCiISKkhu9WtRUpFyEtlPRLxVMLW1p0iuTOkO5HUlt7cTWytgqF/hyQaxWRWdEZ7SFiymBYsMPE/IVkW/GHxHCF5KIc01CMNlkP+f+QVld89kkn5Ddz272+Zg5M+TkfD7vs4ezO28OZ8/HY4wxAgAAABARXrc7AAAAAAxmJNwAAABABJFwAwAAABFEwg0AAABEEAk3AAAAEEEk3AAAAEAEkXADAAAAETQk2gEty9KpU6eUmpoqj8cT7fAAAAD4EmOMWltblZ2dLa83ttZj29vb5ff7w/4+OTlZKSkpUeyRc1FPuE+dOqWcnJxohwUAAEAv6uvrNW7cOLe7EdTe3q6J40eqoSkQtk1WVpbq6upiOumOesKdmpoqSZrygzVKSo7+wIxoCP8XFmmfzHDvX4wevzv/m9A56TNX4kqS1TbUtdjp/xn1t1aQ/8YW12L7dqe6FjtQ0OxK3LPNw1yJK0k3fu0912L/4/3LXYt9+cbwK12RduYXna7FnjL6Y9div/XmFFfiJn3m3v+EZ77l3jxrXNoR9ZjWZx2q+/HaYJ4WK/x+vxqaAnr/rRylpXbPo1paLV1+bb38fj8J9xdd2EaSlJyiJF/0B2bIUPcSbm+Kewm31+vOh1ZguHElriTJci/hTkp2L+FOGh79D+pgbBf+ER00vN2VsF6/e685eWSya7G9w9x73UOS3PssTRruXuyhI1z8+3YpkUmy3Eu4hwxxc56597pjdbvvyFSPRqZ275ul2Ozvl7mXFQAAAAB90GkC6jTdF/E6jeVCb5wj4QYAAEBM65Iluw1dXSLhBgAAAC6aJSNL3Ve47epiEQk3AAAAYlqnMWG2lJBwAwAAABfNb4z8Nsm1XV0sIuEGAABATLP+Vezq4wEJNwAAAGJal/Go03Q/ArDLpi4WkXADAAAgpgXkUcDmzG27ulhEwg0AAICY1mm86jTdH0bUGR9buNWvxyitW7dOEyZMUEpKivLy8nTw4MGB7hcAAAAgSfIrKWyJB44T7m3btqmkpESlpaWqrq5Wbm6u5s6dq6ampkj0DwAAAAnOGI8sm2LiZA+344R77dq1Wr58uQoLCzVt2jRt2LBBw4cP15YtWyLRPwAAACQ4v0kKW+KBo4Tb7/erqqpKBQUFn9/A61VBQYEqKyttr+no6FBLS0tIAQAAAPqqU151Ksmm9Gt3dNQ56uWZM2cUCASUmZkZUp+ZmamGhgbba8rKypSenh4sOTk5/e8tAAAAEk7AeMOWeBDxXq5evVrNzc3BUl9fH+mQAAAAGES6bFe3k9QVJ1+adHQsYEZGhpKSktTY2BhS39jYqKysLNtrfD6ffD5f/3sIAACAhNZphqjTZr+23cNwYpGjFe7k5GTNmjVL5eXlwTrLslReXq78/PwB7xwAAAAQMJ6wJR44fvBNSUmJli1bpmuvvVazZ8/W448/rra2NhUWFkaifwAAAEhw8b7C7TjhXrx4sT7++GOtWbNGDQ0Nmjlzpnbs2NHti5QAAADAQOiU1/YIwE7Fx6Mm+/Vo9+LiYhUXFw90XwAAAIBuLHll2eyEtquLRf1KuAEAAIBo6TRJGmK7pWQQr3ADAAAA0eI3Q5Rkuqet/vjIt0m4AQAAENss45Fl8wVJu7pYRMINAACAmNZlhqjTZoW7ixVuAAAA4OJ1miQlsYcbAAAAiIyApIC6bx8JRL8r/ULCDQAAgJjWaQ1RktU9be20WOEGAAAALlqXSbJ90mSXsVzojXOuJdy33/NP+UYOjXrcfb/8b1GPecHYN92bFJ3/8xNX4uZecsqVuJJ09KEZrsX+rznu/V1f8b86XIvdcItroVU0ZY8rcR8/6t6LTh3a7lrsn1+3y7XYj90z37XYr1z9uGuxG7pSXYv9P779H67E/eVfF7sSV5KOz++e3EXLrZd9GPWY/rN+fRD1qH3HKSUAAABABHWaJHltvzTJCjcAAABw0Ui4AQAAgAhiSwkAAAAQQV1hVrj50iQAAAAwALosr7yWTcJtxcdJ3CTcAAAAiGmWPLJsHnxjVxeLSLgBAAAQ0zqtJHlsVrg7bepiEQk3AAAAYlpAXnUZr219PCDhBgAAQEzjlBIAAAAggrrCbCnpYksJAAAAcPG6jFcemy0ldttMYpHjXu7du1cLFixQdna2PB6PXnnllQh0CwAAADjvwpYSuxIPHCfcbW1tys3N1bp16yLRHwAAACBEl+UNW+KB4y0l8+bN07x58/rcvqOjQx0dHcGfW1panIYEAABAAgsYj+2WksBgXeF2qqysTOnp6cGSk5MT6ZAAAAAYRBJuS4lTq1evVnNzc7DU19dHOiQAAAAGkYHeUrJu3TpNmDBBKSkpysvL08GDB/t03QsvvCCPx6NFixY5ihfxhNvn8yktLS2kAAAAAH0VMN6wxalt27appKREpaWlqq6uVm5urubOnaumpqYer/voo4/0wAMP6MYbb3QcMz52mgMAACBhDeSWkrVr12r58uUqLCzUtGnTtGHDBg0fPlxbtmwJe00gENCSJUv00EMPadKkSY5jknADAAAgpgUsb9ginT+U44vliwd2fJHf71dVVZUKCgqCdV6vVwUFBaqsrAwb/9e//rXGjBmj++67r1/9d5xwnz17VjU1NaqpqZEk1dXVqaamRidOnOhXBwAAAICeWGGSbetfCXdOTk7IIR1lZWW29zlz5owCgYAyMzND6jMzM9XQ0GB7zb59+7R582Zt2rSp3/13fCzgW2+9pZtvvjn4c0lJiSRp2bJl2rp1a787AgAAANgxkoyxr5ek+vr6kO8J+ny+AYnb2tqqe+65R5s2bVJGRka/7+M44Z4zZ46M3SsGAAAAIiBgvJLtOdzn6/p6MEdGRoaSkpLU2NgYUt/Y2KisrKxu7T/44AN99NFHWrBgQbDOsixJ0pAhQ1RbW6vJkyf3Gpc93AAAAIhpAcsTtjiRnJysWbNmqby8PFhnWZbKy8uVn5/frf3UqVN15MiR4Hbqmpoa3XHHHbr55ptVU1PT5+fLOF7hBgAAAKLJGI+MzYkkdnW9KSkp0bJly3Tttddq9uzZevzxx9XW1qbCwkJJ0tKlS3XppZeqrKxMKSkpmj59esj1o0aNkqRu9T0h4QYAAEBMC1heyeYhN4F+PPhm8eLF+vjjj7VmzRo1NDRo5syZ2rFjR/CLlCdOnJDXO7CbQEi4AQAAENMsS/LYbB/513Zqx4qLi1VcXGz7u4qKih6v7c8hISTcAAAAiGkDuaXEDSTcAAAAiGmW8chjk1z350mTbiDhBgAAQGyzPDJ2J5I4PKXELSTcAAAAiGnGhHnwTZw8Gsa1hPuSoWeVMjT64R97cl3UY16w9N9XuhY7fesYV+JWZnQ/RD5azuW79y4c9Z+uhdaf33jatdjLZn/XtdiPzSlwJe68y991Ja4kTR122rXYz5zIcy22vO69t/d/Nsm12GOGtLgW+67U/+dK3KpvHHAlriSVZPzDtdjfe3dp1GN2tXVEPaYTxvLK2JxIYlcXi1jhBgAAQEwz1vliVx8PSLgBAAAQ0zilBAAAAIggY+y/NEnCDQAAAAwE4zlf7OrjAAk3AAAAYpv5V7GrjwMk3AAAAIhtlsf+zG3O4QYAAAAuHqeUAAAAAJHEHm4AAAAgcjzW+WJXHw9IuAEAABDb2MMNAAAARFCcn1Li6AH0ZWVluu6665SamqoxY8Zo0aJFqq2tjVTfAAAAgM9XuO1KHHCUcO/Zs0dFRUXav3+/du3apc7OTt12221qa2uLVP8AAACQ4C7s4bYr8cDRlpIdO3aE/Lx161aNGTNGVVVVuummm2yv6ejoUEdHR/DnlpaWfnQTAAAAiE+OVri/rLm5WZI0evTosG3KysqUnp4eLDk5ORcTEgAAAAnGYzzyWDYlTo4F7HfCbVmWVq5cqRtuuEHTp08P22716tVqbm4Olvr6+v6GBAAAQCKyeihxoN+nlBQVFemdd97Rvn37emzn8/nk8/n6GwYAAAAJzmPOF7v6eNCvhLu4uFivvfaa9u7dq3Hjxg10nwAAAIDPhVvNHowr3MYY/exnP9P27dtVUVGhiRMnRqpfAAAAgCQF92zb1ccDRwl3UVGRnnvuOb366qtKTU1VQ0ODJCk9PV3Dhg2LSAcBAACQ4BLpwTfr169Xc3Oz5syZo7FjxwbLtm3bItU/AAAAJLiEOofbmDj5ZwQAAAAGj3DJ9WBMuAEAAICoi/MtJSTcAAAAiGkJeSwgAAAAEDWscAMAAACR4zH2e7hZ4QYAAAAGQLgTSQblKSUAAABA1LGlBAAAAIgcVrj7aUP1N+QdlhL1uFve/lbUY15wbprftdidue7EHjm8w5W4kvTfx/yXa7F3H7rKtdivnp3sWuwbdta5Frv+w6+4EnfXR1NciStJ1nj3Hml8usGd8ZakIa2Ontk2oN49l+1a7H3+y12L/dvaLFfi+oZ0uRJXkr71f/Jdi/3kA09GPWZbq6V5UY/adyTcAAAAQCSxpQQAAACIHFa4AQAAgAgi4QYAAAAiiS0lAAAAQOTwaHcAAAAggthSAgAAAEQSW0oAAACAyPGYMCvccZJwu/f0AAAAAKAPLmwpsSv9sW7dOk2YMEEpKSnKy8vTwYMHw7Z9+eWXde2112rUqFEaMWKEZs6cqWeeecZRPBJuAAAAxDbTQ3Fo27ZtKikpUWlpqaqrq5Wbm6u5c+eqqanJtv3o0aP1y1/+UpWVlXr77bdVWFiowsJC7dy5s88xSbgBAAAQ0wZyhXvt2rVavny5CgsLNW3aNG3YsEHDhw/Xli1bbNvPmTNH3/72t3XllVdq8uTJuv/++zVjxgzt27evzzEdJdzr16/XjBkzlJaWprS0NOXn5+tvf/ubk1sAAAAAjvSWcLe0tISUjo4O2/v4/X5VVVWpoKAgWOf1elVQUKDKyspe+2GMUXl5uWpra3XTTTf1uf+OEu5x48bpkUceUVVVld566y3dcsstWrhwoY4ePerkNgAAAEDf9bKlJCcnR+np6cFSVlZme5szZ84oEAgoMzMzpD4zM1MNDQ1hwzc3N2vkyJFKTk7W/Pnz9cQTT+ib3/xmn7vv6JSSBQsWhPz829/+VuvXr9f+/ft11VVXObkVAAAA0Ccey8hjdd+wfaGuvr5eaWlpwXqfzzeg8VNTU1VTU6OzZ8+qvLxcJSUlmjRpkubMmdOn6/t9LGAgENBLL72ktrY25efnh23X0dERsqzf0tLS35AAAABIQL09+ObCdufeZGRkKCkpSY2NjSH1jY2NysrKCnud1+vV5ZdfLkmaOXOmjh07prKysj4n3I6/NHnkyBGNHDlSPp9PP/nJT7R9+3ZNmzYtbPuysrKQJf6cnBynIQEAAJDALjza3a44kZycrFmzZqm8vDxYZ1mWysvLe1xA/jLLssLuE7fjOOGeMmWKampqdODAAa1YsULLli3Tu+++G7b96tWr1dzcHCz19fVOQwIAACCBDeQpJSUlJdq0aZOefvppHTt2TCtWrFBbW5sKCwslSUuXLtXq1auD7cvKyrRr1y59+OGHOnbsmB577DE988wz+v73v9/nmI63lCQnJweX1GfNmqVDhw7pj3/8ozZu3Gjb3ufzDfg+GgAAACSQME+a7M853IsXL9bHH3+sNWvWqKGhQTNnztSOHTuCX6Q8ceKEvN7P16Tb2tr005/+VCdPntSwYcM0depUPfvss1q8eHGfY170o92dLqkDAAAAjhhzvtjV90NxcbGKi4ttf1dRURHy88MPP6yHH364X3EucJRwr169WvPmzdNll12m1tZWPffcc6qoqHD0pB0AAADAid6+NBnrHCXcTU1NWrp0qU6fPq309HTNmDFDO3fudHQOIQAAAOCEJyB5bL556AlEvy/94Sjh3rx5c6T6AQAAANgKdyKJ01NK3HLRe7gBAACASOrtwTexjoQbAAAAMS2h9nADAAAAUTfAp5REGwk3AAAAYhor3AAAAEAEeQJGHq/NHu4AK9wAAADAxTOyf6pkfOTbJNwAAACIbR4T5pQS9nADAAAAF4893P10/6w3NGxk9MOvP7ww6jGDOm0ekRQl5vgId+KeGOlKXEk6ua7WtdiZ/+ZxLfa/H/q2a7EveeWoa7FTX+xwJe700addiStJtc2ZrsU2He59nl2ef9y12Ds/vNK12PdO3e9a7P+d839diXvtP5e7EleSvv+Tv7sWu/BAYdRjWufaJf026nH7ymOM7Wo2K9wAAADAAPAEjDw2j5XkS5MAAADAQLDM+WJXHwdIuAEAABDTPOZ8sauPByTcAAAAiGlsKQEAAAAiiS0lAAAAQORwSgkAAAAQSZaR7LaPsMINAAAAXDyPZeSxecqN3dMnYxEJNwAAAGKbMeeLXX0cIOEGAABATPMEjDzilBIAAAAgMixLstlSIsumLgZ5L+biRx55RB6PRytXrhyg7gAAAABfcmFLiV2JA/1e4T506JA2btyoGTNmDGR/AAAAgBDxvqWkXyvcZ8+e1ZIlS7Rp0yZ95StfGeg+AQAAAJ8LWOFLHOhXwl1UVKT58+eroKCg17YdHR1qaWkJKQAAAECfJdqWkhdeeEHV1dU6dOhQn9qXlZXpoYcectwxAAAAQJJkLPsvSJpBuMJdX1+v+++/X3/5y1+UkpLSp2tWr16t5ubmYKmvr+9XRwEAAJCgAoHwJQ44WuGuqqpSU1OTrrnmmmBdIBDQ3r179eSTT6qjo0NJSUkh1/h8Pvl8voHpLQAAABJPIj345tZbb9WRI0dC6goLCzV16lStWrWqW7INAAAAXLSAZb99JE7O4XaUcKempmr69OkhdSNGjNAll1zSrR4AAAAYEJaRZJdwD8IVbgAAACDqLEv2CfcgXOG2U1FRMQDdAAAAAMJI9IQbAAAAiCQTCMiY7ieSGGsQnlICAAAARJ0x9vu1B+MpJQAAAEDUBQKSx2Y122bVOxaRcAMAACCmmUBAxibhtttmEotIuAEAABDbjJHElhIAAAAgMgJWmC0lnFICAAAAXLTzW0q83evZUmLP/Gvpv/1sV7RDS5ICHe2uxJUk6zN3XrMkWe1JrsQN+D2uxJWkLtPpWuxAp4vzzHJzzP3uxW7rcCWuPznxXrMkWZ+5N8fdfN2Bc+697vaz7n2mtQxxZxUxUcfbcuF1W5+df1+ZGN2i0Wn5ZWy2lHTJvb8nJzwmyiN78uRJ5eTkRDMkAAAA+qC+vl7jxo1zuxtB7e3tmjhxohoaGsK2ycrKUl1dnVJSUqLYM2einnBblqVTp04pNTVVHo+zlbiWlhbl5OSovr5eaWlpEerh4MKYOceYOceYOceYOceYOceYOZeoY2aMUWtrq7Kzs+X1dt+64ab29nb5/eH/NzE5OTmmk23JhS0lXq/3ov/llJaWllBvgoHAmDnHmDnHmDnHmDnHmDnHmDmXiGOWnp7udhdspaSkxHxC3ZvY+icMAAAAMMiQcAMAAAARFFcJt8/nU2lpqXw+n9tdiRuMmXOMmXOMmXOMmXOMmXOMmXOMGSIh6l+aBAAAABJJXK1wAwAAAPGGhBsAAACIIBJuAAAAIIJIuAEAAIAIIuEGAAAAIijmEu5169ZpwoQJSklJUV5eng4ePNhj+5deeklTp05VSkqKrr76av31r3+NUk/dV1ZWpuuuu06pqakaM2aMFi1apNra2h6v2bp1qzweT0iJ96c3OfGrX/2q2+ufOnVqj9ck8hyTpAkTJnQbM4/Ho6KiItv2iTjH9u7dqwULFig7O1sej0evvPJKyO+NMVqzZo3Gjh2rYcOGqaCgQO+9916v93X6eRhPehqzzs5OrVq1SldffbVGjBih7OxsLV26VKdOnerxnv15f8eT3ubZvffe2+3133777b3eN1HnmSTbzzaPx6NHH3007D0H+zxDZMRUwr1t2zaVlJSotLRU1dXVys3N1dy5c9XU1GTb/p///Kfuvvtu3XfffTp8+LAWLVqkRYsW6Z133olyz92xZ88eFRUVaf/+/dq1a5c6Ozt12223qa2trcfr0tLSdPr06WA5fvx4lHocG6666qqQ179v376wbRN9jknSoUOHQsZr165dkqTvfe97Ya9JtDnW1tam3NxcrVu3zvb3v//97/WnP/1JGzZs0IEDBzRixAjNnTtX7e3tYe/p9PMw3vQ0ZufOnVN1dbUefPBBVVdX6+WXX1Ztba3uuOOOXu/r5P0db3qbZ5J0++23h7z+559/vsd7JvI8kxQyVqdPn9aWLVvk8Xj03e9+t8f7DuZ5hggxMWT27NmmqKgo+HMgEDDZ2dmmrKzMtv2dd95p5s+fH1KXl5dnfvzjH0e0n7GqqanJSDJ79uwJ2+app54y6enp0etUjCktLTW5ubl9bs8c6+7+++83kydPNpZl2f4+0eeYJLN9+/bgz5ZlmaysLPPoo48G6z799FPj8/nM888/H/Y+Tj8P49mXx8zOwYMHjSRz/PjxsG2cvr/jmd2YLVu2zCxcuNDRfZhnoRYuXGhuueWWHtsk0jzDwImZFW6/36+qqioVFBQE67xerwoKClRZWWl7TWVlZUh7SZo7d27Y9oNdc3OzJGn06NE9tjt79qzGjx+vnJwcLVy4UEePHo1G92LGe++9p+zsbE2aNElLlizRiRMnwrZljoXy+/169tln9YMf/EAejydsu0SfY19UV1enhoaGkHmUnp6uvLy8sPOoP5+Hg11zc7M8Ho9GjRrVYzsn7+/BqKKiQmPGjNGUKVO0YsUKffLJJ2HbMs9CNTY26vXXX9d9993Xa9tEn2dwLmYS7jNnzigQCCgzMzOkPjMzUw0NDbbXNDQ0OGo/mFmWpZUrV+qGG27Q9OnTw7abMmWKtmzZoldffVXPPvusLMvS9ddfr5MnT0axt+7Jy8vT1q1btWPHDq1fv151dXW68cYb1draatueORbqlVde0aeffqp77703bJtEn2NfdmGuOJlH/fk8HMza29u1atUq3X333UpLSwvbzun7e7C5/fbb9ec//1nl5eX63e9+pz179mjevHkKBAK27ZlnoZ5++mmlpqbqO9/5To/tEn2eoX+GuN0BDIyioiK98847ve4jy8/PV35+fvDn66+/XldeeaU2btyo3/zmN5HupuvmzZsX/POMGTOUl5en8ePH68UXX+zTqkai27x5s+bNm6fs7OywbRJ9jmFgdXZ26s4775QxRuvXr++xbaK/v++6667gn6+++mrNmDFDkydPVkVFhW699VYXexYftmzZoiVLlvT6Je9En2fon5hZ4c7IyFBSUpIaGxtD6hsbG5WVlWV7TVZWlqP2g1VxcbFee+017d69W+PGjXN07dChQ/X1r39d77//foR6F9tGjRqlK664IuzrZ4597vjx43rjjTf0wx/+0NF1iT7HLswVJ/OoP5+Hg9GFZPv48ePatWtXj6vbdnp7fw92kyZNUkZGRtjXzzz73D/+8Q/V1tY6/nyTmGfom5hJuJOTkzVr1iyVl5cH6yzLUnl5echq2Rfl5+eHtJekXbt2hW0/2BhjVFxcrO3bt+vvf/+7Jk6c6PgegUBAR44c0dixYyPQw9h39uxZffDBB2Fff6LPsS966qmnNGbMGM2fP9/RdYk+xyZOnKisrKyQedTS0qIDBw6EnUf9+TwcbC4k2++9957eeOMNXXLJJY7v0dv7e7A7efKkPvnkk7Cvn3n2uc2bN2vWrFnKzc11fG2izzP0kdvf2vyiF154wfh8PrN161bz7rvvmh/96Edm1KhRpqGhwRhjzD333GN+8YtfBNu/+eabZsiQIeYPf/iDOXbsmCktLTVDhw41R44cceslRNWKFStMenq6qaioMKdPnw6Wc+fOBdt8ecweeughs3PnTvPBBx+Yqqoqc9ddd5mUlBRz9OhRN15C1P385z83FRUVpq6uzrz55pumoKDAZGRkmKamJmMMcyycQCBgLrvsMrNq1apuv2OOGdPa2moOHz5sDh8+bCSZtWvXmsOHDwdP1HjkkUfMqFGjzKuvvmrefvtts3DhQjNx4kTz2WefBe9xyy23mCeeeCL4c2+fh/GupzHz+/3mjjvuMOPGjTM1NTUhn28dHR3Be3x5zHp7f8e7nsastbXVPPDAA6aystLU1dWZN954w1xzzTXma1/7mmlvbw/eg3kW+t40xpjm5mYzfPhws379ett7JNo8Q2TEVMJtjDFPPPGEueyyy0xycrKZPXu22b9/f/B33/jGN8yyZctC2r/44ovmiiuuMMnJyeaqq64yr7/+epR77B5JtuWpp54KtvnymK1cuTI4vpmZmeZb3/qWqa6ujn7nXbJ48WIzduxYk5ycbC699FKzePFi8/777wd/zxyzt3PnTiPJ1NbWdvsdc8yY3bt3274XL4yLZVnmwQcfNJmZmcbn85lbb72121iOHz/elJaWhtT19HkY73oas7q6urCfb7t37w7e48tj1tv7O971NGbnzp0zt912m/nqV79qhg4dasaPH2+WL1/eLXFmnoW+N40xZuPGjWbYsGHm008/tb1Hos0zRIbHGGMiuoQOAAAAJLCY2cMNAAAADEYk3AAAAEAEkXADAAAAEUTCDQAAAEQQCTcAAAAQQSTcAAAAQASRcAMAAAARRMINAAAARBAJNwAAABBBJNwAAABABJFwAwAAABH0/wETcGhOj6EjQQAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 1000x200 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=(10,2)); imshow(scale[:100].reshape(mix.shape));  colorbar();"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vbjax",
   "language": "python",
   "name": "vbjax"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}

as a tl;dr in code

import jax, jax.numpy as np, jax.random as jr
import numpyro, numpyro.infer
from numpyro.distributions import Normal, Bernoulli

key = jr.PRNGKey(42)
n_class = 5
n_obs_dim = 20
n_obs = 40
confusion = 0.1
class_means = jr.normal(key, shape=(n_class, n_obs_dim))
obs_classes = jr.randint(key, shape=(n_obs,), minval=0, maxval=n_class)
obs = class_means[obs_classes] + confusion*jr.normal(key, shape=(n_obs, n_obs_dim))
obs_classes_oh = jax.nn.one_hot(obs_classes, num_classes=n_class).T
    
def log_p():
    # sample parameters
    mix = numpyro.sample('mix', Normal(np.zeros((n_class, n_obs_dim)), 1))
    offset = numpyro.sample('offset', Normal(np.zeros((n_class,1)), 1))
    params = mix, offset
    # forward model
    probs = jax.nn.softmax(mix@obs.T + offset, axis=0)
    # Bayesian loss
    numpyro.sample('obs_classes_oh',
                   Bernoulli(probs=probs),
                   obs=obs_classes_oh)    

guide = numpyro.infer.autoguide.AutoDiagonalNormal(log_p)
optimizer = numpyro.optim.Adam(step_size=1e-3)
svi = numpyro.infer.SVI(log_p, guide, optimizer, 
                        loss=numpyro.infer.Trace_ELBO())
svi_result = svi.run(key, 5000, progress_bar=False)

loc = svi_result.params['auto_loc']
scale = svi_result.params['auto_scale']
z_mix = (loc[:100].reshape(mix.shape) - mix) / scale[:100].reshape(mix.shape)

assert np.percentile(np.abs(z_mix), 50) < 1