plot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "db9dfed1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import torch\n",
    "from torch.distributions import Normal\n",
    "import math\n",
    "import scipy.stats as stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ae878b6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters of data-generating distribution\n",
    "prior_mean = 10; prior_std = 3; obs_std = 2; seq_length = 7\n",
    "\n",
    "def bayesian_inference(inputs, prior_mean, prior_var, noise_var):\n",
    "    \"\"\" Performs Bayesian inference for the given data-generating process.\n",
    "    Args:\n",
    "        inputs: PyTorch Tensor of shape (num_observations, batch_size, 1)\n",
    "        \n",
    "    Returns:\n",
    "        m: PyTorch Tensor of shape (batch_size, 1)\n",
    "        s: Float\n",
    "    \"\"\"\n",
    "    n = inputs.shape[0]\n",
    "    k = 1 / ((1 / prior_var) + (n / noise_var))\n",
    "    m = k * ((inputs.sum(0) / noise_var) + (prior_mean / prior_var))\n",
    "    s = math.sqrt(k + noise_var)\n",
    "    return m, s\n",
    "\n",
    "# sample data\n",
    "mu = Normal(prior_mean * torch.ones(1), prior_std).sample((1000,))\n",
    "x = Normal(mu, obs_std).sample((seq_length,))\n",
    "\n",
    "# run Bayesian inference\n",
    "m, s = bayesian_inference(x[:seq_length-1], prior_mean, prior_std ** 2, obs_std ** 2)\n",
    "bayes_optimal_loss = -Normal(m, s).log_prob(x[-1]).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "e6b4e0ab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAALIAAACeCAYAAACW5FAlAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAAAig0lEQVR4nO2de5wcVZn3v09VV3dP9/TckplMJjdISCQXUJJwZ1cExACKiCi+iuCVZfEjrouiIquv6+4qsu+KgKuw3hUVBXRZlosBQUCFJEAI4ZILSchtksllZnp6ZvpSXc/7R1X31HRmkkmYnkl36vv59Kerq06dPtX9q1PPOec5zxFVJSCg0jHGuwABAaNBIOSAqiAQckBVEAg5oCoIhBxQFQRCDqgKyi5kETFF5HkRuX+IYyIit4jIehFZJSILy12egOpkLGrkzwCvDHPsPGC297oS+N4YlCegCimrkEVkKnAB8INhkrwb+Jm6PA00iMjkcpYpoDopd418M3Ad4AxzfAqwxfd5q7cvIOCgKJuQReSdQIeqPru/ZEPs22fMXESuFJEVIrJi8pw5uui2WzSVyaiXNnhV1qsslLNGPh24UEQ2Ab8GzhKRX5Sk2QpM832eCmwvzUhV71DVxaq62DKCjpaAfSmbKlT1S6o6VVWPAj4A/FFVLytJdh9wudd7cQrQrartB8579MsbUNmExvoLReQqAFX9PvAAcD6wHugDPjqSPLR8T6iACmVMhKyqjwOPe9vf9+1X4FMHnR+QzuVIZjK01dWNUikDKpmKNThzOlxHSMCRSMUJWQFVxc4HQg4YoOKEXKA/lxvvIgQcRlSkkFXdWjkgoEBlChnF8YScDmrmACpSyK6Ic46DqrK3v39Q7ew4DrtSqXEsX8B4UIFCdk2LTN4m77gNvt5stnjMdhxyTtAQPNKoSCEDZO08/XYOx3FIZjLs7esb7yIFjCMVKGTBMg0e2/AaqUyWbD6P4zikbRs7nw9q4yOUMR+ifqOETZN4OMxvXnwREeHd8+ahQATo6O0d7+IFjBMVVyPHLItb33khi6dO4d6XVrM71UtfNovtNf4CjkwqTsgAhgjvnbeAvKP8z6uvYDsO/baNoxqI+Qil4oQsIlimSVtdHWfMmMFjGzawM5UiY9uk7VyxfzngyKLyhAzEw2GskMnF8xcgIvxi5fPYjkNfNkfWtnEch+3J5HgXNWAMqTghg9vgqwlZTIzHuewtb+GFHe3cvfpFbMehN5cbGCwJuuSOGMrWayEiUeAJ3A6FEHC3qn61JM2ZwH8DG71d96rqP48gbyKmiWOFOPPomWzq7OSBNWv48+uvc9r06Vy+cBGmCGnbRlURGWpqYEA1Uc7utwxwlqqmRMQCnhKRB71p/36eVNV3HkzGk2pr2ZlKEQ1Z5FW54oSFTKuv55mtW3h43TounDuPkCHERdjd20tzbe2oXVTA4Uk55+ypqhacHizvNWotsXg4jGkYxEIW0bDFWbOO4dLj3oyi/HXLZtI5m5w3XN2dTo/W1wYcppQ7QIspIiuBDmCpqj4zRLJTReQFEXlQROaPKF+gNhwmZBiETJNYyCJsmsxsbKQtUccTGze6jb9cjrzj0JvNYufzo3hlAYcbZRWyquZV9S240/xPEpEFJUmeA2ao6puBW4HfD5WPP65F1969GIaBaRi0eCaDZZrUWBbhUIi3Hj2TjZ17WbVjB1nbpt/OoaqkfI5FAdXHmPRaqGoX7uTTJSX7kwXzQ1UfACwRmTjE+cW4Fs3NzYOOTaqtRUSIhkJEQiZnz5rFhFiM365eRWe6n/5sjmw+T18uF9TKVUw5Iw01i0iDt10DnAO8WpKmVbwuBRE5ySvPnoP5HtMwaKur83oyQiSiET62cDHtPT3809Kl7O3vpy/neskVfDH6gtq56ihnjTwZeExEVgHLcW3k+0XkqkJsC+ASYLWIvADcAnxAD3GMubGmhpDXv7xw6hRueNtZdKfT3PPSarK2TSafR1VxHIeuoPFXdZSt+01VVwEnDLHfH9fiNuC20fi+GsuiP+faw7GQxTFNE1gyZw6PrF/HadNncHxrK5ZhsCOYPVKVVOTI3nA0xWKICGHTJBa2uHjefCbG4/zX8mV09veTDWzkqqWqhAyuiWEYBtGQRX1NDX934sns7e/jqdc3kc7bgXdclVJ1Qja94WjTMIhZFrMnTmRCLMZre/Zg5x2ynq2cD2aSVBVVJ+QCk2prsQwDyzSZ29zC6o6d5PJ5Ml6tHIz2VRdVK2TTMDC8Wvm4Sa0k0xk2dnaSydnY3hy/gOqh6oQcDoWYEIsVP1uGwaIpbRgi3PfKy6QyGXJOYF5UG1UnZIBIyO1VbInHMQyDSbUJPnXyKaxsb+fXL66iP+eaFzuDrriqoSqFXCBkmq5zkWly1jGzWDx1Cs+3byeXzwddcVVGVQsZoC4aBSAasjhpyjS6+vtZt2c3mUDIVUXVCxncEAKmCCdOnYqI8Oy2beS8brjA76I6OCKE3OANkkyIxVjYNoUnN22iO50ml88HfhdVwhEhZIBEJIJlmrx3/gJ6c1keXreWdD7ogqsWjhghxywL0zCY0zyR41sn8+hrr5HJBUKuFo4YIZveQpNhw+T06TPozWZZs3s3qko2GBypeI4YIYPrUBQ2TRZPnYppCMu2biGXz7M7iH9R8RxRQq6xLAzDoD4a4U0Tm1mxbStZJ+iGqwZGJGQRiYuI4W3PEZELvVgV+zsnKiLLvBnSL4nI14ZIIyJyi4isF5FVIrLw0C5j5CQiESzDZF5LCx2p3iAaUZUw0hr5CSAqIlOAR3GX2v3JAc4pBGh5M/AWYIm33rSf84DZ3utK4HsjLM8hE7csQobBUY2NhAyDDXv24jhOMDG1whmpkEVV+4CLgVtV9T3AvP2dMMIALe8GfualfRpoEJHJIy/+oWEaBrOaJpBXh01dneR8E1MDKpMRC1lETgU+BPyvt++A8/1GEKBlCrDF93mrt680n2Jci127do2wyPstF821ceqjUTZ1dgZecFXASIX8D8CXgN+p6ksiMhN47EAnjSBAy1DRBfeZi7S/uBYHi+F1w0XMENPrG9jc3Y0drGtd8YxIyKr6J1W9UFVv9Bp9u1X1mpF+yXABWnBr4Gm+z1OB7SPN91CZnEgU7eQdqZ5gOeAqYKS9Fr8UkToRiQMvA2tE5PMHOOeAAVqA+4DLvd6LU4BuVW0/2Is4WEQEEWFGQwOOKtu6k8Gk1ApnpKbFPFVNAhcBDwDTgQ8f4JyRBGh5ANgArAf+C7j6IMv/hjiqsQk777C5u4u84+AEtnLFMtIALZbXb3wRcJuq5kRkv1XYCAO0KPCpkRd3dJlRX48hsD2ZJOc42I5D2DiixoiqhpH+a7cDm4A48ISIzAAqepGOxpoa4pEIzfFatnZ342gQR7mSGVGNrKq34MZmK/C6iLytPEUaGyKmScgwmFpfz85UqhgUPKAyGWljr15E/qPQlysi/w+3dq5YCt1w0xsa2JHqIZ0Nei4qmZGaFj8CeoD3e68k8ONyFWqsEBGObmwkY+dd8yKokSuWkQp5lqp+VVU3eK+vATPLWbCx4riWSQCs7tiJQtANV6GMVMj9InJG4YOInA70l6dIY8eEWIxJiQQttXFe7ujAdhzae3rGu1gBh8BIu9+uAn4mIvXe507givIUaeyIhEKEDIP5LZN4ZssWsnm7GNwloLIY6RD1C5475vHA8ap6AnBWWUs2RogI81ta6MvlWLfnoFZ9CDiMOKjef2/xmkL/8T+WoTxjjojwppYWDIHX9gZCrlTeyDBW1ayLOzmRwDJNNnd2BY29CuWNCLkq/vFEJELEMJlSV8/rXV04qqQymfEuVsBBst+WjYj0MLRgBagpS4nGmEQkwp7eXqbV1/P89u3FlVJrI5HxLlrAQbDfGllVE6paN8QroapV07wPGQbT6utJZjLs6e0lH5gXFUfg6oXb4Jve0ADA+s6941uYgEMiEDKFoeomTEN4vbMTVWVvXx/pXI7tyQEnv+50OggfcJhSziV8p4nIYyLyihfX4jNDpDlTRLpFZKX3+kq5ynMgmmI1RMwQ25JJHFX6czn29ruDl47jkM7l6M1mSdv2kJNVAz+N8aWcdq4NXKuqz4lIAnhWRJaq6ssl6Z5U1XeWsRwHRADLMGlNJNje00MunydsmogIqrrPaqld6X5ChokpUlxYx1GlPholHg6TtW3C+xkhtPN5QqYJuDeAETjzv2HK9guqaruqPudt9wCvMMRU/8OBVm8y6uREgvZkkkw+X2zw6RDvGTtPbzZLMpNxF2z3juUdh/5cblAsucL61346enuLAcZ3pFLkDiE4TDqXC4KU+xiTqkBEjsKd9lQa1wLgVC+s1oMiMn8sylOKiGAaBm2JOvb299OTyRTn8Dne6k+quo+oCxSOpbJZOj1zJGvbZG2b9p4edqRSbE8mydh2cRaKMmCO+JdKy+XzdIxgkZ6udDoIUu6j7F1oIlIL3AP8g294u8BzwAxVTYnI+cDvccNnleZxJW5ILaZPn16Wchoi3gifwbbubuqiERTX7BARd0VVVXcG9jB5+AU+VITPPSU19Y5UatA5juOuzGoPY2/3ZDJk8zYTYm9sToOq0pfLEQ+HR5xe5MADuduTSVpra+lKp4uB1ceKsgrZm7B6D3Cnqt5betwvbFV9QET+U0QmquruknR3AHcALF68uCydvCLCtPp6VGF7Tw/HTJhA3lBMwxWu6cZwJGQYrph9onZKBF6wrYf7HlUl6Rs97Mlk6PE+F/Io9JbURSIkMxkmJxLFNKU1di7vmjoxy8LybPvtySRh0yyuXhWzLBpqarDz+WJ4MFOEbD5PjbfGCrimTnPcvVFsxyk+YSYnEogIWdtmd18fbXV1AGS8p0nBa9BRJW3bWKaJZZp0pFK01NYe1H9xKJRNyOLewj8EXlHV/xgmTSuwU1VVRE7CNXXGzXNnWmMDoGzYu4fTpk9HVMk7giFCThwEIa9OUdQi7jFHtSgEcGt3fw1WWpcNZaIUxe+dVzhWMEXae3qKNaO/xvZ3D/b5As2o6qAl2PpyuUHHgWKvTMqztQvXsGuIOHilftr+7/XT631Hob1gOw55xykGWi8X5ayRT8eNffGiF/8N4HrcmBiFsACXAH8vIjauo/4HdJy8dgwR4iGLE9qm8NfNm3nfcccTNk13hVQYqH0dQcTx7OoBiaoYKIogKFoUeyFvvDwMKP6xjqprugCi6poyvsf4cLZ4IcDM/n4q/zFVLfaMDHVO4fvyvvyHS1fY7y+jf7s3my2aLoUbp7AwZ2ttbdl6aMomZFV9igN4yKnqbcBt5SrDwdCaSLAxl+PcY2azsn07D61dwztmzyk+MtUTXcGMMERwnAFHFNtwioJVBUMUEfcHKOw3kGJ6xRW04E6EFW9fIX2puYIndBgQ+3C2q1+MgitQvFpcfccL1+N/gmiJmAtlKhwb6rtKtx1VDNinfDtSqaJJMtpUjb/EaGCKcGxzM2+Z3MbvXn6J1Tt38qW3nkltJOLWNLiP+vpodJ8/1VEpNgzVEwcM1FyGV3s7jvvnWii5vIMIWGhRYIYIguCgGAwIWb28HHXc/SJFkQ0Sm/du+MqRdxy0UCND8UlgeHkouA1Z3/nqON736T61kYgUj4N7Y8HATeGogvfEwf9blPFhGwjZhyGCZZp85ayz+cP6dXz/maf5z2eepsay+OzpZ/BKRwdfXvow15x6Oie0tQ06t/jI9f2p4vV0ABjqCquwv9C1B+AYWqz9CoJHB24CETe/kGFgO66tbpnutv8pUEgLDBKlo0oIX/eht98siM679sITQ70by0Sw1fcd3nWaPjOq8F2F4wI4eF2W3rbp3Xgj6fk4VAIh+zAMg3DIpDud5l3HzuXR9etZsW0rAO889lgeWLsGgFU72jmhrY2U11PQmkgUexL8dU7pY1cHPmD7HttaeOyr4ui+tmdBlI6jxRpSvf7tgolQFJlXExZq4sITIi/OoLIJYBsyqMBuLe9rgOLmYRRsJC3Y21oUsu9e9a7fJ3AZeAKFDCneEOUgEHIJlmECOXL5PP/89rezp6+P6x58gF88/zwvdewkZJg8u30bOcfhL6+/jgicM+sYrjr5FEKG4e3fxNzmFhKRCF946EFOaGvjwrkDAf4HiZ3Bgl+3ezdburt469EzMRgQcSFdsdYr9EiUNM7UG7xxCmaBrxYeZDuLQH7wo95fi/tvDqfkZjFKG6EM9LeXmiLFm1WNgadNGQiEXELIMDBE6OzvpzWRoD4S5bQZR7F03Tqa43Heu2ABty97hj+/volLFhxHTybjrqJq2+zq7S2+jm1pJmyabOrqZLMnzPpolPpolPaeJE9u2sRZM2dhOw43/+XPzGpq4uL5C/jBimVsSyZZvnUrF7zpWOZPmlQsm5a8F/cP0eAqvCczGfb29zGjoXHI436GErKXeMCWhqI5Qkn6QQ1SoD+XI2pZxRvPonwDJIGQfURDIVQVyzTJ2jY7vL7T82bPoSEa5dLjjscQ4cG1a1gy+028Z/58N75yMskTmzYyt7mFmU1NvGvuXO5ZvZpM3uaDb34Lv3vpJW5Y+gcm1dbSHI+Ttm1Wtm+ntTbBq7t3sdZ7rduzm/aeHs48eiYvd3Rw05NP8I13LGFeSwuRUIiMbZPKZnhu+3Y2d3Uxvb6Bs2bNQnGHxHtzWfLOgC3dk8nw9ccepSPVyzfesYTJiQTgBnAsDHT4KTQEo1aIrJ2nPholbJrFfmW/PZ22baI+x6jCub9Y+Tzbkkla4nH+unkzN513Po017mSicvasSqVNtly8eLGuWLGiLHmrKu09bgT7vmy2OPAw1C8UNk3qolFChkF/Lkd3Ok2rJ5RSXtyxg0fWr6MznWbVjnYynghMw6A3m+V9C46jKRbj9mXPEAmF+OX7P4CIcOXv7yVmWSxsm8IrXgCZjZ17B9WIx7W28ubWydz14io+d8bfsrCtjXg4jKPKl5c+zMsdHZgiLGybwg1vOwsRYW9/H6lMllU72gkZJqfPmEEiEuGuVav4+crnmFJXx/Vnvo0ZDY209/SQtW1aamv58bMrePS19cxtaeGF9h18+4J3FlfHUlV+uGI597y02vVTUTAN4YoTFnLxguN4YuMGjp3YzCkzZpTFvqg8ITc364r3vGfwzre/Hd73Pkin4ZohVoR417vcV1cXXHfdvscvuQTOPRd27qT3C18o9vEWhNx58cX0nHwy1tattN5yyz6nd192GeYpp9C7ejUzfvhD0rZNxDSLjZveK6/EXrCA0OrVOLfdxt6+XkKGyaod7bTU1jL7376BzpnDU7/5Dcf/7/1Mq28A3H7X57dv4yfnLaFhzps4bu1azn1mGa2JBDWWRXsyydrdu7l5yTvobWrklFde4SMbN5HKZsnaebrS/RzX2sojn/gkd6xfy0e3t3PaqlVs7uoaZFffcdmHOGHmLPJ33cW7tmwlmclgGgbHt7ayfOtWfvWpq7npvPP59ec/x6yXXS9cQwya4zFmtk1h3Vf+ibgV5sHrruPczk7a6uro7O9nWzKJ0djAxi9+ke8ve4Zrnnuej//kp2URcmBalFDoMjNFUK+mCZkGlidMf2OpIIRcPk93Xx8RXPdKGPBBAOjq76fgbZAIh0mEwyhurR6zLDKOgwm8bdYsYvUNxSdAa20t58w6htMvupi+tjaisTgNr20o5ju9oYGW2lp6Fi1i1rz5/HHLVjZ2dlITsjBEmNHQyNS6ei6aPx+jqZFNP/85m7u6mFJXT8gwaIhGqQmH+bEZ4g/r1vL5tiksUtezbtmWLTy9ZTOmYbBm9y7+7fHHmJhJc2xzMxNicTpSKdbt2c2WjRu48f7/IWSYXCQws6mJeDjMhFiMnJPnpVSKH6xYQWttomiqleV/q7gauYymRYGCH4G/Vna8ftG8DsRRVmegcePv6ipwoF82ZlnFYdyGmhq6SuzW+miU7nSaqBUinXNvDP93FEyZZDpNXy7HxHicnT09tNXVDRqpy+TzWF4jdqgZ4jt7esg6+eKTAGBj515+sHw5F86dy+MbNvD4RvcG+ukl76eltpa+bJafPPcsb2puRhDuevEFTp0+gyWz59BSW4shwvo9u/nKI0s5YXIbnzzpJL739NPceuG7A9MCxkbIecdhp8/FctBAhg74J+fRor1aGKhwHX/c/lN/37HfR2GfoV/fdwsHvgEORI1lkbbtYRtX4ZBJ1t7XmT9mWURCIQwR9vT1EQmFaKypYUtXFzc+8Sdsx+H/nn0OzfE4jmoxjaMO9dEa98mUTlNjWcTDYXZ7jcQJsRimYdDd38/81tZAyDA2QgbozWYHLcUwnE9B0deiZJaHrY7X7+rW4uoNDPgHFsSQYq1eGHgo3Cyl3+kf+nYO8J8N6hJjsIPPSG+a0mMZbzpXzLIO+UazDIPjJk8ObOSxJB4OEw+HSWUyJDOZwW6ZJb4DBSciGBCfhVnsknIcBzU9YXk2eLFP1nRrddsRLM/Xudhbou5IGeoOXTve8HDecQY7WeiAnwMy2OQB91x1FPXdDMWBlaEcgfA5LnmfY5ZVvFb//tL3Yh4lDkMjdc4/VAIhH4DaSITaSISsbRcFa3hi8/vtmiVukkUTwnOgGW4QozBaFlr1PE3f+hb51lawbbo/9zlyM2YM+GB4/hkighZmXqjrnuk4Do43NFxwFsozIDZTDJziE0IJiVH0oSiO3in4R6wFqHnoYXrf/nYwDBru+jW9Z5+NtrSQcxxMEfI64LZaEL5q8V71fgeIP/wHsG36Lrhg1P+fAoGQR0jprGjLNGmJxwl5PssFx/NCrVNjhXDUnahaYJAZV1I7hUyT7Lnn0vuJTxBavZra++8n9elPD7KlGcaXt9C74vdl8Nv0ruiMQaZGqFAz+26o0lq18ZFHcJYsAdMke9mHKUyMMsUhVPBu85eDAQ++vFdriwgRw0BNk3wZpz6Vc4bINOBnQCvgAHeo6ndK0gjwHeB8oA/4iHozryuBwpR+ERnWz7azv4/+nI3l+WGAO4Lon3AKJc73fX0QjxO/807Cy5cj2Sw9116LNjURv+02er76Vcjnqb/2Wrpvvpno/fcTeeghAHqvuYZ8ayuNN9wAItgzZ9L7mc8M8mGOLl1K9N57wTBIffaz5GfPpv6qq7BnzcJau5bej3wEp7GR8Guv0fz5z9N/wQWEn32Wvg9/GGv1aqxlyzB6e8E0yZx8MpFHHyU/dSqp664jtH49tbfeCtksmdNOo++yywh5N6BdiUJmZHEtzsOdbDobOBn4nvdeNTTWxGisGZhYahpun3QynSaVzdIQjZJXJRqNEnrsMRpefRU2b6bv298m1dJC9vLLadqzh6bbb0e//nXI5QgDoZdfRk47jbjjoMuXk//Rj8h0dRH72teIfvCDpBcvInn5FcVnvIiQiETo6e+n5p576PrudzF27aL2O98h+c1vYnR303/55fTW1dHwuc/R9d3vYh9zDN3f/jaYJpHnngNx3Ue1sZHkV79K7U03YeRyJG+9lfprr8Xo6SE/fTrd3/kOiFD/mc/Q//73I4aBlDwxRptyzhBpB9q97R4RKcS18Av53cDPvOlNT4tIg4hM1jFYj3qsMQyDqM80SEQixMPhgblsoRCcfz5cfTXs3Uvs618n9jd/Aw8+6JohIm6ac86Bp56CFSvg4x+HbduQdesIXX31wJ954onUvvACtTfeSP9JJxGeOxfz3/8dJk4kdeWVGG1tTGpoYGcohKRSTIzFCE+YgHH00e5E23B4kB/F5EQCiURIJBIQieLMnUt9XR1MmwYLFqChEFZrK5MBenrg5pshnUa3bqWnq4v6SIRkX1/lN/b2E9diCrDF93mrt6/qhFxKMcTAUMRi0NsLv/0t3HknbN0K//Iv7rGzz4YbbgDbhilT3LTz5sG3vuUet213WtNV7nLfNR/8IPzyl3DHHQA0ZDKEOzowHIfJPT1IU5N7gySTRPbuhbo6UKUpFiMXiVBvWfsI0PCbCCI0xWJgmm7tf/fdcMUVsGgR8vGP01pbC6ZJY00NuVhsVH9DP+Md12Kof3Kf/qCxiGtxWPDAA7ByJWSzbm37pz/BJz8JCxcOpInHIRyGk05yPzc2whlnuOkMA0480U3/3e+6oi6k86iJRODSS+ETn0AMA77wBfdAQwPcfjusXevmBVhvfSvW9dfDRReN/BrOOANuvBFmzgSvyw7cG3d/YcTeMIUuonK8AAt4GPjHYY7fDvwf3+c1wOT95blo0SI94rn+etVdu0Y3z499bHTzG56yaK2c0TiFA8S1AO4DLheXU4BurUL7eFT513+FpiaYOHG8S3JYUbYhanEXmHwSeBG3+w1K4lp4Yr8NWILb/fZRVd3v+PNYDVEHlI3KGqLWkcW1UOBT5SpDwJFDEJg3oCoIhBxQFQRCDqgKAiEHVAWBkAOqgkDIAVVBIOSAqqDi5uyJuz72mvEuRxmZCOw+YKrKJaqqC0Y700qcIbJGVRePdyHKhYisqPbrK0e+gWkRUBUEQg6oCipRyHeMdwHKTHB9h0DFNfYCAoaiEmvkgIB9qCghi8gSEVkjIutF5IvjXZ6RIiKbRORFEVlZaLWLSJOILBWRdd57oy/9l7xrXCMi7/DtX+Tls15EbpFyzubc//X8SEQ6RGS1b9+oXY+IRETkLm//M96cz/1Trqkno/0CTOA1YCYQBl4A5o13uUZY9k3AxJJ93wK+6G1/EbjR257nXVsEONq7ZtM7tgw4FdfP+0HgvHG6nr8FFgKry3E9wNXA973tDwB3HahMlVQjnwSsV9UNqpoFfo0bTqBSeTfwU2/7p8BFvv2/VtWMqm4E1gMnichkoE5V/6ruP/wz3zljiqo+Aewt2T2a1+PP627g7AM9fSpJyMOFDqgEFPiDiDzrzQgHmKTe/ETvvcXbP9x1TvG2S/cfLozm9RTPUVUb6AYm7O/LK2lkb0ShAw5TTlfV7SLSAiwVkVf3k3a466zU6z+U6znoa62kGnkrMM33eSqwfZzKclCo6nbvvQP4Ha6ZtNN7vOK9d3jJh7vOrd526f7DhdG8nuI5IhIC6tnXlBlEJQl5OTBbRI4WkTBuI+C+cS7TARGRuBf7DhGJA+cCq3HLfoWX7Argv73t+4APeC33o3Hj4i3zHtc9InKKZy9e7jvncGA0r8ef1yXAHz07enjGu0V/kK3l84G1uC3fL493eUZY5pm4rfYXgJcK5ca1+R4F1nnvTb5zvuxd4xp8PRPAYtyb4DXcMAoyTtf0K9ywZjnc2vPjo3k9QBT4LW7DcBkw80BlCkb2AqqCSjItAgKGJRByQFUQCDmgKgiEHFAVBEIOqAoCIb9BRCTlvR8lIh8c5byvL/n8l9HMv5oIhDx6HAUclJBF5EDLHA0SsqqedpBlOmIIhDx6fBP4G8/n+LMiYorITSKyXERWicjfAYjImSLymIj8Ejd2NCLye8+h6KWCU5GIfBOo8fK709tXqP3Fy3u15897qS/vx0XkbhF5VUTuHC+f5TFnvEe+Kv0FpLz3M4H7ffuvBG7wtiPAClx/3DOBXuBoX9om770Gd6Rrgj/vIb7rvcBSXB/tScBmYLKXdzeu34IB/BU4Y7x/o7F4BTVy+TgXd1mJlbirWU3A9TMA19dgoy/tNSLyAvA0rrPMbPbPGcCvVDWvqjuBPwEn+vLeqqoOsBLX5Kl6KsmNs9IQ4NOq+vCgnSJn4tbI/s/nAKeqap+IPI7ra3CgvIcj49vOc4T8x0GNPHr0AAnf54eBvxcRC0BE5njeb6XUA52eiI8FTvEdyxXOL+EJ4FLPDm/GnXq0bFSuokI5Iu7WMWIVYHsmwk9w19g+CnjOa3DtYuipSQ8BV4nIKlzvsKd9x+4AVonIc6r6Id/+3+HOdXsB1+H8OlXd4d0IRySB91tAVRCYFgFVQSDkgKogEHJAVRAIOaAqCIQcUBUEQg6oCgIhB1QFgZADqoL/DynoyFFjfY/2AAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 144x144 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.rcParams[\"figure.figsize\"] = (2.,2)\n",
    "\n",
    "losses = torch.load('losses.pth')\n",
    "losses_mean = losses.numpy().mean(0)\n",
    "losses_std = losses.numpy().std(0) \n",
    "\n",
    "plt.plot(torch.arange(losses.shape[1])[::100] , losses_mean[::100], color='#007977', alpha=0.8)\n",
    "plt.fill_between(torch.arange(losses.shape[1])[::1], losses_mean + losses_std, losses_mean - losses_std, color='#007977', alpha=0.1)\n",
    "sns.despine()\n",
    "\n",
    "plt.axhline(y=bayes_optimal_loss, xmin=0, xmax=losses.shape[1], ls='--', c='red', alpha=0.8)\n",
    "plt.text(4750, 1.98, 'Bayes-optimal', c='red', alpha=0.8, size=8)\n",
    "\n",
    "plt.xlim(0, 10000)\n",
    "plt.ylim(1.75, 4)\n",
    "\n",
    "plt.xlabel('Iteration')\n",
    "plt.ylabel('Loss')\n",
    "\n",
    "plt.savefig('loss.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "93090e84",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MetaLearner(torch.nn.Module):\n",
    "    def __init__(self, num_hidden=256):\n",
    "        super(MetaLearner, self).__init__()\n",
    "        \n",
    "        self.rnn = torch.nn.LSTM(1, num_hidden)\n",
    "        self.linear_mu = torch.nn.Linear(num_hidden, 1)\n",
    "        self.linear_logstd = torch.nn.Linear(num_hidden, 1)\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            inputs (Tensor): PyTorch tensor of size (seq_length, batch_size, 1)\n",
    "        \n",
    "        Returns:\n",
    "            Normal object: Normal distribution object with dimensions (batch_size, 1)\n",
    "        \"\"\"\n",
    "        hidden, _ = self.rnn(inputs)\n",
    "        hidden = hidden[-1]\n",
    "        \n",
    "        mu = self.linear_mu(hidden)\n",
    "        std = torch.exp(self.linear_logstd(hidden))\n",
    "        \n",
    "        return Normal(mu, std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "af022a5b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAKMAAACtCAYAAADcZs/wAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAAAfvklEQVR4nO2deXxV5bX3vysDhCFkYg6EMEtAwxAVVCCiRhCHUq0XtAqtvEC5Dryt9wPXSxVa+xbrvda22guKSCu10gqoVKikzIqKhNEAgRAGwxSSEAKE4SRZ7x97n3AIOSc7yZly2N/PZ3/OOXt49jonK89+hvX8lqgqNjbBQFigDbCxcWI7o03QYDujTdBgO6NN0GA7o03QYDujTdBgO6OPEZFYEZnqpbKaishiEckVka9FJNkb5QYLtjP6nljAK84IPAWcVtUewG+BV7xUblBgO6PvmQN0F5HtIvJqA8t6CPiT+f5D4C4RkQaWGTREBNqA64AZQD9V7V/TQRHZCETXcOh5Vf1XtX2JwHcAqlouImeABKDQe+YGDtsZA4yqDq3D6TXVgiEzn2s7Y4CpY82YD3QG8kUkAogBin1sot+wndH3nKVmZwPqXDN+AowHvgQeAdZoCEW6SAh9l6BFRN4HbgJWqup/NKCcKOA9YABGjThWVfO8Y2XgsZ3RJmiwh3ZsggbbGW2CBtsZbYIG2xltggbbGW2CBtsZbYIG2xltggbbGW2CBtsZbYIG2xltgoZG54wjR45UjLCp63rLzMzUzMzMgNtRj80tjS5qp7AwJOJIG8zFixcDbYLXaXQ1Y2Nh4cKFls9NT0/3mR3VWbhwYZ1s8ye2M/oIf/zBKysrfX4Pf9LoHtOBYt26dbz++uuUl5dTWVnJ6NGjWbRoEffccw8ZGRm88MILlJeX8+KLL1JRUcH27dtJT09n9uzZbNy4kczMTC5fvsyiRYvo3r17jfdYunQpr732GgC///3vGThwIA8//DCFhYVER0ezbNkyjh49ylNPPUVZWRlDhgzh1VdfJTU1lS+++II5c+YwcuRI3nzzTRYvXkxERAQLFiygY8eOPProo1y8eJHY2FhGjhzpz5/OMiHjjNP+sZztx481qIz+HTry+v0PuD0eGRnJRx99xOTJkxERvvzySwYPHszmzZtZuXIlTZo04f777+ezzz6jf//+rFu3DoCbb76ZmTNnsnbtWubPn8+vf/3ra8qurKzkzTffZP369ZSWljJx4kSWLFnCe++9R/PmzXnppZdYs2YNvXv3prCwkPnz5xMeHs6KFSuYPXs2ly5d4umnn2bQoEGsWbOGDRs2kJOTw29+8xvS09MZMmQI06dPZ9KkSQ36jXxJyDijP0hJSQGgQ4cO9OnTB4CoqCh27NjB6NGjgZo7WAsWLGDx4sU4HA769OnDrl27eOaZZ0hKSuLPf/4zAKdOnSI7O5u77roLgIiICMrLy5k2bRp79+6loKCAbt260bt3b1JTU7n55psBaNeuHa1btwagpKSEvLw8srKyqtqhnTp1Ii8vj4EDBwIwaNAgH/06DSdknNFTjeYtXJcou74fMGAAS5YsoVmzZjgcjmuOL1iwgKysLFavXs2iRYu48cYbq2pNJ61btyY1NZWVK1cSFhaGw+Fg+/btqCobNmxg5syZOKPyw8KuNPVd76OqdO3alcGDB/PBBx8A4HA4WLp0KTt27CAjI4Nt27YxePBg7/0oXiRknDGQzJw5k1GjRiEipKSk8OabbzJw4EDGjBnD9OnT6d+/P+np6aSmprotIzw8nJ/85Cekp6cTHh7OiBEjmDZtGtnZ2dx3333ExMTQo0ePqvNXrlxZYzlt27YlPT2dYcOGER4ezrhx4xg/fjw/+MEPyMzMJC4uzuvf31vUugZGRP4beFdVs/1jkmfS0tJ0y5YtgTYj4CxfvhyABx7w/RPBy7hVwLAytLMXeMsUGpoiIjHes8vG5gq1OqOqzlfV24EngWRgp4i8LyJ3+to4m+sLS4PeIhIO3GBuhcAO4Kci8kEt140UkRxTwm2Gm3PSTVGkbBFZX0f7bUKIWjswIvIa8ACwBvh/qrrZPPSKiOR4uC4ceBO4B0OW4xsR+URVd7ucEwv8ERipqkdEpG29v8l1RlJSUqBN8DpWetPfAjNVtayGY7d4uO4WINepeGDWog8Bu13OeQxYqqpHAFS1wJLVNh575o0VK4/px6s7ooisBlDVMx6uq5JvM8k397nSC4gTkXUikiUiT1qwJyCsW7eO5ORk0tPTuf3229mzZ4/Xyp4zZw4nTpzwWnmNFbfOKCJRIhIPtBaROBGJN7dkoKOFsq3It0UAg4DRwL3Az0WkVw22TBKRLSKy5dSpUxZu7RsmTJjAunXrePXVV5k7d67Xyp0xYwbt27ev0zXLly+vGt4JFTzVjJOBLIxOy1bzfRbwMUZbsDac8m1OOgHVJ4/zgX+q6nlVLQQ2ANc8f1T1LVVNU9W0Nm3aWLi1bzl79izR0dE8/PDDDB8+nPvvvx+Hw8GcOXNYtWoVAIsWLeKdd95hz5493HXXXdxxxx0sWLAAgCeffJJhw4YxbNgwKisrmTBhAocOHSIrK4vhw4eTlpbGvHnzAJg1axaTJ0/m7rvvZuLEiQH7zv7AbZtRVX8H/E5EnlHVP9Sj7G+AniLSFTgKjMVoI7ryMfCGqTXYBLgVQ6u67qycASd21evSKtrfCKPmuD28cOFC1qxZQ05ODuvXrycpKemqIIZx48bx8ssvk5GRwbJly3j77beZMmUKixYton379owcOZInnniCw4cPs2HDBlT1qum8lJQU1q9fT0VFBcOGDWPy5MkA9OvXj3nz5pGRkUFJSQmxsbEN+55BiltnFJERqroGOCoi369+XFWXeirYlPl9GvgMCAcWqGq2iEwxj89V1T0i8k9gJ1AJzFfVbxvwfXzKhAkTmDVrFtnZ2fzyl7+kRYsWVwUx3HvvvZw4cYKCggJUlfj4ePbv38+4ceMAI4iiqKiIiRMn8vjjj9OtWzdmz55dVX5ubi7PP/88Fy9eZPfuK/28fv36AdCxY0fOnDlz/TkjMBxjOKem+SYFPDojgKquAFZU2ze32udXgYYKr3us0bxNTEwMx44do0ePHtcEMdx7771MnTqVMWPGANCrVy/eeOMN2rRpg8PhICwsjLFjx/LEE08wadIktm3bVlXuvHnzePHFF7n11lvp2bNn1f7qwRChiqfH9Evm64/8Z05ws3DhQtatW8fly5d55ZVXmD59+jVBDI8++igzZsyoah/Onj2bxx9/HIfDQXx8PO+88w4PPvggFRUVxMXFVYWlAdx3331MmTKF/v37ExPjeda1W7duvvuigUJVPW7Ac0ArjN7xfIzOTEZt1/lqGzRokAYzBQUFOnbs2ECbEcy4/dtaGWf8saqWAhlAW+BHGLlNbKqRnZ3NQw89xLPPPuvze5WXl1NeXu7z+/gTKzMwzgbLfRihZDtCKRGON+nbty+bNm3yy72c8YyNMITMLVZqxiwRWYXhjJ+JSDRGz9fGxqtYqRmfAvoDeapaJiIJGI9qGxuvUqszqmqliJwEUszBaRsbn2AlhOwV4N8wom0qzN2KMXVnY+M1rNR03wN6q+olH9tiUwd69bomnqTRY8UZ84BIwHbGIKJ3796BNsHrWHHGMmC7GcNY5ZCq6vvBNBu3OFXIoqKiAmyJ97DijJ+Ym00QkZmZCYTWOKOV3vSfRKQZkKSqbte82Ng0lFoHvUXkAWA78E/zc38RsWtKG69jZQZmFsbiqhIAVd0OdPWZRTbXLVacsVyvXXgVukF1NgHD0lJVEXkMCBeRnsCzgH+iAWzc4hoHGSpYqRmfAfpiDOv8FSgFpvnQJhsLdO/e3a0CbmPFSm+6DPgvc7MJEs6dOwdAy5YtA2yJ9/BYM4rIeBHZKiLnzW1LMC+0v55Yu3Yta9euDbQZXsXT6sAnMR7HP8VYaiDAQOBVEUFV/+wXC22uGzzVjFOBMaq6VlXPqGqJGktXHzaPhTSOigomLVvK2L++z9lL9rS8P/DkjK1U9VD1nea+VlYKtyKJZ553s4hUiMgjVsr1B7/ZsJ63v9nM4l07+dmKTwNtznWBJ2e8UM9jwFWSeKOAFGCciFwzHmGe9wrGYv+gQFWZ+/VXpLbvQEbPnvxl+zbO2bWjz/HUm+4jIjtr2C+AlUW7ViTxwBg6WgLcbKFMv7D92DHyS0sZP3AQHVpFs2r/fv51IJfvpfQNtGlV3HTTTYE2wet4dMYGll2TJN6trieISCIwBhhBEDnjpzl7ARianEyrqCiaRkTw6d69QeWMXbp0CbQJXseTosThBpZtRRLvdWC6qlZ4Wv0qIpOASeAfxdbPDx8msVUrWrdoAUCfNm34Jj/f5/etCyUlJQAhpbvjy0SWViTx0oAPROQQ8AjwRxH5XvWC1I+SeKpK1tF8UtpeUXS+oU0b9pwqwFFR4eFK/7Jx40Y2btwYaDO8ii+dsUoST0SaYEjiXRV6pqpdVTVZVZOBD4GpqvqRD22qlWOlpRSWldGv3RXxzpS27bhcUUH2yZMBtCz0sRLPeL+I1NlpVbUccEri7QH+pqYknlMWLxjZeeI4AL3NfHwAfdsZteTm/O9qvMbGO1hxsrHAfhH5jYjUqVOjqitUtZeqdlfVX5n75mo1WTxz/wRV/bAu5fuCHccNZ+yRkFC1r3NMLM0iI4Ou3RhqWElK9ENgAHAAeFdEvjQ1tqN9bl0A2HXyJLFRUcQ2a1a1L0yErnFx7Dppi8D7EkuPX1OFbAnwAdABYzhmq4g840PbAsLugpN0qaGHmhQbS25Rkf8NcsPAgQOr0vaGClbajA+KyDIMFdtI4BZVHYUhBP+8j+3zK6rKvsJCusXHX3Osa2wcRWVlnDGXiAaaxMREEhOrZzJp3FiJ9H4E+K2qXiVnYopA/dg3ZgWG/DNnKHM46B6fcM2xbgmGg+4tKODWIMhOVWTW0gkJ19raWLHymD5e3RFN/R1UdbVPrAoQ+4sKAUiuISdz1zjDGXcXBMfwzqZNm/ymBekvrDjjPTXsG+VtQ4KBHDPhUeca2oyJrVoRJsLuAjujnK/wFFz7E4y4xe7VAiaigS98bVggyC0qIlyEdjWE8keGh9O+ZUtyCgsDYNn1gac24/vASuDXgGss4llVLfapVQEit7iIti1bEhFW8wOjU0xw9ahDDU+PaTUDaf8dOOuyYeYUDDnyiovpGO0+brhzTAyHSk6HdC6WQFJbzXg/Rr5A5eooHMVaTGOj4nBJCSM85FdJio3lgsPByXPnaB8d2DH/W27xlF25ceIphOx+8/W6kDI5faGMs5cu0Tkm1u05yXHGsQPFRQF3xnbt2gX0/r7AUwfG4/C+qm71vjmBI7fQaAvW1JN2kmQ66t5Tp7i9S7LvjfLASTOCKJSc0tNj+n88HFOM6OyQYZ/ZS07ykCatQ3Q0YSLsC2DOayebN28GrhN9RlW905+GBJp95oB3Yiv3HZiI8HDatmjJfrtH7RNqTfFbU3pfqD3Fb2Mjt6iIuKhmNG/SxON5HVtFk1dsO6Mv8GmK38bEweJi2kfXrlvTKSaGtXl51yQut2k4dopfk0Mlp+nbtvbOQOeYGM5eukTxhTISmrfwg2XXD1aSEiUALwF3YNSInwO/UNWQeVZdKi/nxNmz3OuScNwdXWKNIIrcoqKAOuNtt90WsHv7CiuBEh8ApzA0dh4x3y/2pVH+5mBxMQo1BtVWJ9k8Z9+pwM5RJyQkhFT4GFiLZ4xX1V+6fH65puWkjRln77hz7LWhY9VJjIlBgH2FgR3eOXr0qGFPCAXYWqkZ14rIWBEJM7dHgZBSQtpvjjF28jCs46RpRARtWrRgX4CHd7Zu3crWrSE17+BxaOcsV+akfwosMg+FAecw2pEhwf6iQqIiIkho3tzS+R1ateKAPbzjddzWjKoaraqtzNcwVY0wtzBV9Yoknog8LiI7zW2TiKQ25MvUlwPFRbRv2dLyUE1idCsOFodkFF1AsZQ/WkTigJ5AVaK66ksRarjGKYl3D4bUyTci8omquqqQHQSGq+ppERkFvEU1cSh/kFdcTGIr99OA1ekcG8OKfRcouXDhqiWtNg3DyurAiRi5pT8DZpuvsyyUXSWJp6qXMXrlD7meoKqbVPW0+fErDD0ev1JRWcmRkjN09jAnXZ0ks6NjP6q9i5UOzHMYcnWHzfnqARjDO7VRkySep67fUxiR5X7luzMlOCor6GKGh1mhq7lg60BR4B7VQ4cOZejQoQG7vy+w8pi+qKoXRQQRaaqqe0XESrJjK5J4xokid2I44x1ujvtMEs8ZrZMcZz143VmL5gRweCeUpPCcWHHGfBGJBT4CMkXkNNdK29V4HbVL4iEiNwHzgVHuZnVU9S2M9iRpaWlejfnfa4aDJdXhj9ssMpKE5s2rHNmnXDoL2/4CB9fDhWJoFg/dR3A4fihENgsp0VArSYnGmG9nichaIAYzw2otVEniAUcxBKQecz1BRJIwAi6eUNV9dTHcW+ScOkXT8HDatKjb1F6H6GjfLs5Sha3vwar/gkul0KINNE+AwlzIWcHO8NHQ+Ra6jH8OQiRgw2pveiBX5qa/MDskHlHVchFxSuKFAwucknjm8bnAi0AChkgoGEkz0+r1TepJblERHcw10XWhY3QrtpuKZV6nwgEfTYVdf4O2KTDk36Gdi4TziZ2waj0c2gif7IcHfg9h4b6xxY9YCZR4EfgBV0LG3hWRv6vqy7Vdq6orgBXV9s11eT8RmFgni71MbnGRx4Bad3SOiWFV7n7OX75Mi1piIOtEZQV88BjsXwX9H4O0p651tPY3QXIFFOyBbW/BhdPw6HuN3iGt9KbHATer6ktmWNlg4HHfmuUfyisqOFJS4nGpgTucbcwD3nxUq8Ly5wxHvHUK3DLJvYNJmFFb3jwR9n4KH081rm/EWHHGQ7gMdgNNMbQaGz1HzpRQXllZNW5YF7qaSmVOfR6vsOFV2PYe3PgIpI61ds2AH8JN/wY7PoCv/ug9WwKAp7npP2C0ES8B2SKSaX6+ByOmsdHj1NbpWoPQU204a9Mcb4WS7c+Etb+C5GEwuPZseHcOchldu3UyFOfBqp9Dh/6QfLt3bPIzntqMW8zXLGCZy/51PrPGz3xrLvesSY+xNlo2bUpMVBT7irww1ni+CJZNgZhOMOIF4xFc2/2buzysJAzuehGWToLFP4SpX0J0e/cXBymelh38yfnezFbQy/yYo6oOXxvmD3aeOEFMVBTxFqN1qtPRG8M7qrBsMlw8DSPnQkRU7dcAB/INNbTuncwUIU2jIeNX8NFkwymf/LjRDflYmZtOB/ZjBD38EdgnIsN8a5Z/yD55sipyuz4ktoppeAfmm/mQm2n0mlvXvuzBye6Dx9l9sNrQUkI3o9NzcD18fY2Gf9BjpQPzP0CGqg5X1WHAvcBvfWuW76msrCSn8FRVR6Q+dI2P48S5c/WXVj61D1bNNIZqrHZYaqPfI9BxAGS+ZJTfiLDijJGqmuP8YM6URPrOJP9w5EwJZQ4HPRNa136yG1LaGI/I7ceszI5Wo8IBH/7IGLoZMdNSO9ESInDnf0F4pFF+ReNpUVn5BbJE5B0RSTe3tzE6NY2anceNNBquyYfqSh8zpduWo/XID7P6F3DyWxj6M2jZtvbz60KL1jD0p0b5a37l3bJ9iBVnnAJkA89ihJPtNvc1ajbnf4cAPRrgjAnNmxMbFUWWuTjKMgc/h01/gB53Q3cfSRZ1HwHd7oRNv4MjX/vmHl7G43SgmaYtS1X7Aa/5xyT/8HX+d3SOjaVlA6byRITu8QnsOFGHOeoLJbB0IrRsB3f8tN73vueWa/LIX8uwnxm145KnjOGepsGdR8pjzaiqlcAOM7omZFBVth49Sh8vZGjt1aY1+woLueiw0DZThY9+AucK4K6fQ5P6DSkBRDWNJKppLU33Ji0h/QUozTemGYMcK4/pDhgzMKtF5BPn5mvDfMnhktMUX7hwVebU+jKwQ0fKKyv5+jsLSS63LICcFTBo/NVROPUg5/AJcg5bSB+XOABSx8G3S4y4yCDGSgjZbJ9b4Wc2HTbyuqe2b7gzDkjsCMDavAMM9yDBTMEe+Od/QvtUYz65gew7Yswe9e5i4Tuk/RiObYcVP4POt0LrHg2+vy9wWzOKSJSITMMIH7sBI45xvXPzl4G+IDM3l2YREfRqQOfFSWxUM7rExrLh0EH3J104De8/CpFRcPeL3hvGsUpYBNw9CyTcsONiqX/vbxFPv8qfgDRgF0YSIk9Kto0GVeVfubmkduhIRLh34v9uat+ezd/lU15Rce3BinL44HEoPQr3/NKI1g4ELdsaDnn6IPx9vBE3GWR4csYUVf2hqs7DEHwKiaVoB4qLyC89w21eXNg1JKkL5x2X2Xjo0NUHVI1H4+Ev4LbnoMNNXrtnveiUZkQEHVgDn70QWFtqwJMzVnUPVbXcD7b4hU927wHgjuRkr5U5NDmZJuHhLN6148pOVWOqL2sh9P0+pDzotfs1iBsfgRtGG3PXa2oN1vcrnjowqSLibFwI0Mz8LBgJi+oeqx8ELNq+leS4uDot2q+NZpGRDEpM5OPde/jfh9RYo5v5Enz5BvQaBbd5Py33qCH96n/x0J+B44IRzCvhcOd/es+wBuBJayfc1Npx6u1EuLxvlI64p6CAbcePM7r3DV6XQM7o0ZMT586SufdbYzHVpt9BzwwY/h8+CeWKiAgnIqKebV4JM+avu6XD+jmwfFpQzGH7uVsXWH77xeeEifBgnz5eLzujZ09uaFpO4tLHYMf7cOOjkD7DZz3n7LxjZOfVI0DDSVg4jPg59B0DWe/Cu6OMwfgAct04Y15xEe9mbWFUr96Wpe8sU1lO5z3vkyXv0eXyMXIGTIMhU306hJN39BR5RxsYZR4WDrc/Z/zTHN8BfxgEX78dsJ62T53RgiSeiMjvzeM7a8vKVV/KKyqY8OHfCRfh6SGDvVauVFwmft9HpCwZQ+evXqE8tit3h43joZxILpQ3oj5fr5EwZh7EJsHK52Hu7bBjsd8f3ZYW8dcHi5J4ozCk9npiSOH9L16WxLvocDBhyd/ZeOgQLwxPp02L2tNreELKLxF9fDMxh1YTdzCTyEslXG7Rjvy0ZzmbeDtPHivk6S++ZfTyDSy77w5imnpxTbUvie8KD/7BWCab9S4sm2SoWfR5AG54wFjkFdHUpyb4zBlxkcQDEBGnJJ6rMz4E/FmNnLlfiUisiHRQVbdhMKqVXL5UdlWaXSOeA1BFgaKyMo6UnGbDwYO8m/UN+aWlTBkwkEd6dwXHefPcqqurypHKcsId5wm/fI4wxznCL58l8vwpmp47StPS72hWvI+oM4cRraAyLJKy1n053m0U59ulVj2S70xsw0uDevGLrH30fO9T/k/f7ozo1I5esdHENo2kRWREndUr/IYI9LoXet4DR76Cb5fCtkXGnHpYBCT0MObUY5KMxWMxiYbsSmRzI+gjsoXxGhZp/B4SZpRp8fuKr3Ini8gjwEhTNQIReQK4VVWfdjnnH8AcVf3c/LwamK6qW2oqEyCtY7humdSw2q0+XG7Siost2lPWoiOlsb0ojeuJhruv9XaXXOCN/QVknb5wlfSaABEiVRpt4twEBKn6XBsZYUY9sqrSt82BZjgYzlGGyDFukkL6Ukx7zhNWh/+nSoVKhEqEJrNPu73SlzWjFUk8S7J5rpJ4wDmZXZpT/Rwf0hoohFKM1kbDUFxmExrAstpPsYr5/WrmPLDE3LzCL+SfqjqypkO+dEYrkniWZPNcJfH8jYhs8bcYlT8Jpu/ny950lSSeue56LFA9DvIT4EmzVz0YOOOpvWgT2visZrQoibcCuA/IBcqA6yJPoU3N+KwDEyqIyCSzmRCSBNP3s53RJmi4bqYDbYIf2xk9UNt0ZmNDRBaISIGIfOuyL15EMkVkv/lad31AL2E7oxtcpjNHASnAOBGxsFg5qFkIVB/jmwGsVtWewGrzc0CwndE9tWb4amyYKfaqZ1J6CGO9E+br9/xpkyu2M7qnrhm+GivtnGO75quXhX+sYzujeyxn+LLxDrYzusfSVGUIcFJEOgCYrwEL97ad0T1WpjNDgU+A8eb78cDHgTLEHvT2gIjcB7zOlenMxiN2WAMi8lcgHSNS5yTwEkZOyL8BScAR4AeqGpB0sbYz2gQN9mPaJmiwndEmaLCd0SZosJ3RJmiwndEmaLCd0QUROeen+7jVoxORQyLScBVT9+VPEJGO/rpfXbCdMTAEUhxxAtCxtpMCge2MNWAmX1onIh+KyF4R+YuYsmUiMkdEdptyLP9t7msjIktE5Btzu93c31JE3hWRXeb5D4vIHAx5we0iYknx3UP5s8wYxXUikiciz7pc83PT9kwR+auIPG+uZU8D/mLev5l5+jMistW08wbv/ZJ1RFXtzdyAc+ZrOnAGYz46DPgSuAOIB3K4MlkQa76+D9xhvk8C9pjvXwFedyk/zvU+bmw4BLSuts9d+bOATRgJ6VsDRRip9NKA7UAzIBojEenz5jXrgLRq93vGfD8VmB+o39+X66YbO5tVNR9ARLYDycBXwEVgvoh8CvzDPPduIMVF87GViESb+6syVKrq6Xra4q58gE9V9RJwSUQKgHYY/zgfq+oF0/7ltZS/1HzNAr5fTxsbjO2M7rnk8r4CiFBj+e0twF0YTvY0MAKj9hzi/OM7MR/t3phvdVd+jXZiTSHFFWcZzusDgt1mrAMi0hKIUdUVwDSgv3loFYZjOs9zt9+5vsQhInXJTOuufHd8DjwgRvqUlsBol2NnMR7dQYftjHUjGviHiOwE1gP/19z/LJBmdlJcE32+DMSJyLcisgO409z/FrDTQwdmp4jkm9trHsqvEVX9BiM0bAfGI3gLRhsYjHUwc6t1YIICO2onRBGRlqp6TkSaAxuASaq6NdB2ecJuM4Yub5mrGaOAPwW7I4JdM9oEEXab0SZosJ3RJmiwndEmaLCd0SZosJ3RJmiwndEmaPj/WX3suzit92sAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 144x144 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAIEAAACtCAYAAACXzhwdAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAAAP4klEQVR4nO2deXgUZZ7HP7/OQRKSkINDCGc4HBMIEAE5RTxGRnRwPFYd79F11fWaVffR1XGd2XXGXXfmcQ9dRFYFRBQVdaKyEljQCCMhcoOCGIEcJIGGjnQgnaPf/aMLjZiEpFPd1Ul+n+epp6rfqnrr25VPqt46W4wxKN0bl9MBFOdRCRSVQFEJFFQCBZVAoZNIICIpInK3TXWdKyKbRKRBRK46ZdzNIvKV1d3cpHyYiGywyt8QkVirXETkP0Rkr4hsE5HcJvPMFpHd1rhH7MgeMowxEd8BQ4EdNtaVAywCrmpSngYUW/1UazjVGrcMuNYangfcZQ1fAqwABJgMbLDKo4CvgUwgFtgKZDm9HlvqOsWWAHgaGC4iW0TkmY5UZIzZZ4zZBvhPGXUxkG+MOWKMOQrkA7NFRIDzgbes6RYCl1vDc4FFJsBnQIqI9AcmAXuNMcXGmDrgdWvaiCTa6QBt5BFgtDFmXHMjRaQASGpm1EPGmFVtXEYGUNLkc6lVlg54jDENp5S3Nk9z5ee0MUfY6SwStIoxZoYN1UhzVbdSHuw8EUeXkMCmLUEpcF6TzwOBtcBhApv5aGtrMBAobzLPoFPmKSfQDmiuPDJxulHSxsZcOrDf5jpf4ccNw28INApTreE0a9yb/LBheLc1PIcfNgwLrfJoAg3LYXzfMMx2ej22uC6cDtCOP9prwA7gmQ7WM5HAf3AN4AZ2Nhn3K2Cv1d3apDwTKLTK3wR6WOUCPEfgSGA7MKHJPJcAe6xxjzm9/lrrxAqsdGM6yyGiEkJUAkUlUFQCBZVAQSVQUAkUVAIFlUBBJVBopwSzZ882BC6JdusuPz/f5OfnO56jnV2LtOtS8uHDh9szeZeltrbW6Qi2orsDRSVQVAKFLnJ7WbjJyMg4/USdCJUgCHJzc08/USdCdweKShAMK1asYMWKFU7HsA3dHQRBQ0PD6SfqROiWQFEJFJVAQdsEQTF48GCnI9iKShAEY8eOdTqCrejuQFEJgiEvL4+8vDynY9iGSqCoBIpKoKASKOghYlBkZmY6HcFWVIIgyM7OdjqCrejuIAgaGhq61JVElSAIutr9BCqBohIoKoGCSqCgh4hBMWrUKKcj2IpKEARnnnmm0xFsRXcHQVBbW9ulnkxWCYIgPz+f/Px8p2PYhkqgqASKSqCgEijoIWJQZGVlOR3BVlSCIBg+fLjTEWxFdwdB4PV68Xq9TsewDZUgCNasWcOaNWucjmEbKoGiEigqgYJKoKCHiEGRk5PjdARbUQmCYMiQIU5HsBXdHQSBx+PB4/E4HcM2VIIgKCgooKCgwOkYtqESKCqBohIoqAQKeogYFF3tVfcqQRB0tR+90N1BELjdbtxut9MxbEMlCIL169ezfv16p2PYhkqgqASKSqCgEijoIWJQTJo0yekItqISBEG/fv2cjmArujsIgsrKSiorK52OYRsqQRAUFhZSWFjodAzbUAkUlUBRCRRUAgU9RAyKqVOnOh3BVlSCIEhPT3c6gq3o7iAIysrKKCsrczqGbeiWIAg2bdoEdJ07jHRLoKgEikqgoBIoaMMwKGbMmOF0BFtRCYIgJSXF6Qi2ohK0RslG2PYGVO4EEeg3GsZdx/76NKDrvKxCJWiGkpIvMO/dw+DDRTRINL7kQSREuZCSDVD4Att63QJDpqkEXZEt5eW8+M5zPFm1gJ408PeNk/lvfzbeQ7GclZLEH2aMYO6xdbC5HHa9C+eMhIzxTsfuMGKMafPEEyZMMEVFRSGM4wzGGH63ehWr1y7hf6PzqIlOYvv4B0num4mnrp6Cg0dYtKeEfd4TXJGZwQ1RjUSXbeQyWQ03vANDpjj9FdqCtDSi228JGhobue6NpWzfUcDGHiuQ+DSOznyK/j16AdAzJpprR2RwZWZ/Xty1j+d37adnQgKXDpwC7iJYchXcthL6dd4f0e7W5wka/X6ue2MpBTsKWZfwEfExMZRN+w2NlgBNiXG5uHt0Js/NGMOx+gbeK3FTdf4fwBUNr10DJzzh/wI20W0l8Pv93PrWmyzfsY2ClHWk+msomfwI9T1bv518Zv/e/HzSKD6q83Hh6i85Pusf4dtyWHYT+P1hSm8v3VaCJ1bls3jLZpYNrGDk8b1U5PyK2rSRbZp3ckYffjcli51HvuUXm6ppPOdO+OZj+L9/CnHq0NAtJcj7Yhe/X7uGvx0UyxWH3ufb/hOpHjKrzfNXVHk4KzqGR8ePZGVJJfccHgKZs2Dds7C/8z2y3u0kKHa7uXHZMs5K6ckzNW/S0COZg+PvCpwMaiOl5UcoLT/C9SMHcsOIAczb9Q1L+10NCb3h7dvBdyyE38B+upUEtfX1zH11EQ1+Px/02Umc9yDluffgj00Mus6Hx41kXFoyt63bRXHu/YH2wQcP2pg69HQrCe54Zzk7KitZMjqeoftXcGT4HI73HdOhOmNcLp6dNpo4l3Bx0VHqsq8MnGre9Z5NqUNPt5HghQ2fsXjLZu7JHsacr56jNnkwVdm/tKXuvvE9+NfJWRRX13Bz9RhM6jDIux+8VbbUH2q6hQRFpSXc934euWecwVPHluFqqKVswgOBY3ybmHZGOndmDeH14gqWDrwJfF5Yfge044ysU3R5CY4cP84vXl1Mco8eLB3iJrmiiKrsG6hLHhh0nTnZg8nJHvyj8ruzhzGlbwq3bD7E/p9cB8VroHB+R+KHhS4tgd/v55qlr1Hh9bJgyghGbH0eb9+xHM2c3aF6Y2OiiY358VbEJcIfp4wmNTaGGXvSqOuXAysfh8pdHVpeqOnSEjyxKp9VX+/l4XNymb3lKRpjEimfcF+7Dgebo7ziKOUVR5sdl9Ijhn+fPprK2jquqbsAEx0fOJtYf6JDywwlXVaC17du5am1a7ho+HAe9LxN7LEyys++l8bYpA7X3ZoEAGPTevFY7ijerazlhT5/Be6v4MOHOrzcUNElJSgsOcCtb7/JWX36MK9/Jelff4h75NwOHw62h6szB3B1Zn/uKo5ma/+fwuZXoejlsC2/PXQ5CUo8Hi5dtJBecXEszu3DsI1/xNt3LIeyrg17lsdzRzGxdy8mHRhGVWpWYGuwb13Yc5yOLiVBWXU1M+bPo6aujv85N4exnz5KXc9+lE38O5Dwf9UYl4vnZ+QwLDmJcYencCIuHd64Ho58E/YsrdFlJKg4doxz579AldfLy7Ny+em6B0CgZPKj+GPiHcvVMyaaF2eOw9UjiWk1F1Ff74OXfwaeA45lOpUuIcGeQ4eYOu95yo99y4JZZ3P5hodwNdRyYNoT1CeeYfvyxo8ZyvgxQ9s8fe+4WF6ZNZ7DsX2Y6ZtD3XEPvDQbPCW2ZwuGTi/BR3v2MPH5/8J9/DivTx3G1Rt+TVTDCfZPewJfr6EhWWZUlIuoqPatuoye8bx2wdm4Ewcz03cJPq8bs+ACKP08JBnbQ6eVwNfQwOMrP2LOwpdJjYvjk3FR/HzDwxiJZv/0J/GlDAvZskvK3JSUtf/3DtLjYlk862xcfc7kHN9lVB2vxbx0MXy+MAQp206nvNH04+Kvuf2d5ex1u7lpUDLPsorUzUXU9M6mdNKDHbo03BYqD1UDMCij/W8sSYqNZv65Y3npyxTG7khkWUw+5+bdh3/nO7gu/ROkZdod97R0Ggn8fj8f7tnN79eu4S8HDjA5wcdbQ8vIObgKf1QMldk3cmTEHEeOAtqLS4TbzxrK5DPSeKAonVnV6/ht8afE/edE/OOuJ3b6/ZA+PGx5IlqCb2tr+ezAAZbv2sG7O3eSUHOQa+LKeSW1nFE1u/EfjKY6YypVo2+kMS7F6bjtZnRqMq9eOJH80qFctH0Mf3PiY365aTH+zYuoSB9LUs7lJOVcAamhfdIp6IdPjDFNOj8Yg+GHZd+N+276QJnfGGp8Pry+Wrx1dYG+10OFuxz30Uo81VUc81QS5T3IIPEy2uVhgstNmqkBoC6hL9UDp3M0c7Yjf/yiLcWB9THOvk23MYZNbg8rd+8mq6qAK9nDSAnsdo64EqlIGII3eTCu5Azi0ocSmzqImPhe9EjoRVx8MnHxycTEJeJyReFyRSGuqFOvkbR4waR9EgyINoV/3RNXx66/tIt6icEX3xtf4gBqkgbhSR+NL6FP+AI0w/Y9gUbhmFGheYtZozHsrq7lm4r99PV8ST/fQUb6D5FJNYnS0M66BD9CzG+P2iOBiBwC9rcrRcfoDRwO4/LCSbi/22FjTLPX0NslQbgRkSJjzASnc4SCSPpukd+UVkKOSqBEvASRf4Ne8ETMd4voNoESHiJ9S6CEgYiUQERmi8huEdkrIo84naejiMhLIlIlIjualKWJSL6IfGX1U53KF3ESiEgU8BzwMyALuE5EspxN1WFeAU49Rn8EWG2MGQmstj47QsRJAEwC9hpjio0xdcDrwFyHM3UIY8wnwJFTiucCJ68hLwQuD2empkSiBBlA01tuSq2yrkY/Y8xBAKvf16kgkShBc+e49RAmhESiBKXAoCafBwLlDmUJJZUi0h/A6jv2CHMkSrARGCkiw0QkFrgW+LPDmULBn4GbreGbAedeaPDD+wIiowMuAfYAXwOPOZ3Hhu+zFDgI1BPY0t0GpBM4KvjK6qc5lU/PGCoRuTtQwoxKoKgEikqgoBIoRJAEIuIN03L+oZVx+0SkdwiXfYuIDAjX8tpKxEgQRlqUIAzcAgw43UThJuIkEJHzRGStiLwlIl+KyBKRwFMUIvK0iOwSkW0i8m9WWR8ReVtENlrdNKs8UUReFpHt1vRXisjTQLyIbBGRJW3M01L9T1r3CawVkWIRua/JPL+xsueLyFIReUhErgImAEus5Z98acK9IrLJyvkT+9ZkO3D6bFqTs2peq38eUE3gmoEL+AswHUgDdvP9LXEpVv81YLo1PBj4whr+F+DZJvWnNl1OCxn2Ab1PKWup/ieB9UAPAs8QuIEYAn/oLUA8kETgjOBD1jxrgQmnLO9ea/huYIET6z5Sn0UsNMaUAojIFmAo8BlQCywQkQ+A961pLwSy5PtHrpJFJMkq/+5FRcaYll831jot1Q/wgTHGB/hEpAroR0DY94wxJ6z8eaepf7nV/xy4IsiMHSJSJfA1GW4Eoo0xDSIyCbiAwB/3HuB8AluLKSdX+kmsXYgd58Rbqr/ZnLTyzF8LnKzj5PxhJ+LaBC0hIolAL2PMh8ADwDhr1EoCQpycrqXyk/fw1YtITDsW3VL9LfEpcJmIxFmZ5zQZd4zALiKi6DQSEFh574vINuBj4NdW+X3ABKvxtwu40yr/ZyBVRHaIyFbg5E+bzAe2tdIw3CYipVb3p1bqbxZjzEYCl4m3EtjUFxFo40DgXsN5pzQMHUevIoYAEUk0xnhFJAH4BLjDGLPJ6VwtEaltgs7OfOsO6ThgYSQLALolUOhcbQIlRKgEikqgqAQKKoGCSqAA/w8eEEMFA/YdHgAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 144x144 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def plot_gaussian(mu, sigma, color):\n",
    "    x = np.linspace(mu - 50*sigma, mu + 50*sigma, 1000)\n",
    "    plt.plot(x, stats.norm.pdf(x, mu, sigma), color=color)\n",
    "    plt.fill_between(x, stats.norm.pdf(x, mu, sigma), alpha=0.2, color=color)\n",
    "\n",
    "iters = [0, 10000]\n",
    "\n",
    "# generate test data\n",
    "mu_test = torch.Tensor([[7]])\n",
    "x_test = Normal(mu_test, obs_std).sample((seq_length,))\n",
    "\n",
    "for i in iters:\n",
    "    plt.figure(i)\n",
    "    \n",
    "    # meta-learned inference\n",
    "    network = torch.load('trained_models/iter_' + str(i) + '.pth')\n",
    "    predictive_posterior = network(x_test[:seq_length-1])\n",
    "    \n",
    "    # Bayesian inference\n",
    "    m, s = bayesian_inference(x_test[:seq_length-1], prior_mean, prior_std ** 2, obs_std ** 2)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        plt.rcParams[\"figure.figsize\"] = (2.,2)\n",
    "        plt.title('t = ' + str(i), y=1.05, size=10)\n",
    "        plot_gaussian(predictive_posterior.loc[0, 0], predictive_posterior.scale[0, 0], color='#007977')\n",
    "        plot_gaussian(m[seq_id, 0], s, color='C1')\n",
    "        plt.xlabel('Insect Length')\n",
    "        if i == 0:\n",
    "            plt.ylabel('Probability Density')\n",
    "            plt.legend(['meta-learned', 'Bayesian'], frameon=False, fontsize=7.9, loc=1, bbox_to_anchor=(0.8,1.05))\n",
    "        else:\n",
    "            plt.yticks([])\n",
    "        plt.axvline(x=7, ymin=0, ymax=1, ls='--', c='grey', alpha=0.8)\n",
    "        plt.ylim([0, 0.7])\n",
    "        plt.xlim([-5, 15])\n",
    "        sns.despine()\n",
    "        plt.savefig(str(i) + '.pdf', bbox_inches='tight')\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0866532f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}