notebooks/bayesian_linear_regression_bnn_wrapper.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GJcMJWZDBy2h"
      },
      "source": [
        "# Bayesian Linear Regression VI (using BNN wrapper) vs closed-form solution"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 139,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "The autoreload extension is already loaded. To reload it, use:\n",
            "  %reload_ext autoreload\n"
          ]
        }
      ],
      "source": [
        "%load_ext autoreload\n",
        "%autoreload 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 140,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uJd1yRdoJgA7",
        "outputId": "ac8f8121-db81-47a0-f47e-eb001f0b38c8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Note: you may need to restart the kernel to use updated packages.\n"
          ]
        }
      ],
      "source": [
        "%pip install -q torch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 141,
      "metadata": {
        "id": "PkWZ_bKBJYeT"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "import torch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 142,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'1.13.0+cu117'"
            ]
          },
          "execution_count": 142,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "torch.__version__"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 143,
      "metadata": {},
      "outputs": [],
      "source": [
        "import sys\n",
        "sys.path.append(\"../\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 144,
      "metadata": {},
      "outputs": [],
      "source": [
        "import reparameterized as r\n",
        "from reparameterized import sampling\n",
        "from reparameterized import likelihoods\n",
        "from reparameterized import densities\n",
        "from reparameterized import bnn_wrapper"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 145,
      "metadata": {
        "id": "JVdgdm7Cd5UI"
      },
      "outputs": [],
      "source": [
        "from matplotlib import pyplot as plt\n",
        "plt.rcParams['figure.figsize'] = (5, 5) # Width and height"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 146,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uiEnOgwGJaEn",
        "outputId": "c8936bc9-58b4-4ecf-fcea-88d08592c3c4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using CPU\n"
          ]
        }
      ],
      "source": [
        "if torch.cuda.is_available():\n",
        "    torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
        "    env = torch.cuda\n",
        "    print(\"Using GPU\")\n",
        "else:\n",
        "    torch.set_default_tensor_type('torch.FloatTensor')\n",
        "    env = torch\n",
        "    print(\"Using CPU\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 147,
      "metadata": {},
      "outputs": [],
      "source": [
        "SEED = 123\n",
        "\n",
        "torch.manual_seed(SEED)\n",
        "np.random.seed(SEED)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kTKZ15jMg0nN"
      },
      "source": [
        "## Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 148,
      "metadata": {},
      "outputs": [],
      "source": [
        "# # Toy Data\n",
        "\n",
        "# def generate_toy_data(N):\n",
        "#     x_data = np.float32(np.random.uniform(-10.5, 10.5, (1, N))).T\n",
        "#     r_data = np.float32(np.random.normal(size=(N,1)))\n",
        "#     y_data = np.float32(np.sin(0.75*x_data)*7.0+x_data*0.5+r_data*2.0)+10\n",
        "    \n",
        "#     mask = (x_data<0.0) | (x_data>5.0)\n",
        "#     r_data = r_data[mask]\n",
        "#     y_data = y_data[mask, None]\n",
        "#     x_data = x_data[mask, None]\n",
        "\n",
        "#     return x_data, y_data, r_data\n",
        "\n",
        "\n",
        "# x_data, y_data, r_data = generate_toy_data(100)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 149,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Boston dataset\n",
        "\n",
        "import pandas as pd\n",
        "\n",
        "data_url = \"http://lib.stat.cmu.edu/datasets/boston\"\n",
        "raw_df = pd.read_csv(data_url, sep=\"\\s+\", skiprows=22, header=None)\n",
        "data = np.hstack([raw_df.values[::2, :], raw_df.values[1::2, :2]])\n",
        "target = raw_df.values[1::2, 2]\n",
        "\n",
        "x_data = data\n",
        "y_data = target[:,None]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 150,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 356
        },
        "id": "7pbiMR-uuQw0",
        "outputId": "23ab92bf-fc5f-4e78-ecd3-44c98cd8f6d2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "x=torch.Size([506, 13]) y=torch.Size([506, 1])\n"
          ]
        }
      ],
      "source": [
        "x = torch.tensor(x_data, dtype = torch.float, requires_grad=False)\n",
        "y = torch.tensor(y_data, dtype = torch.float, requires_grad=False)\n",
        "\n",
        "print(\"x=%s y=%s\" % (x.shape, y.shape))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HdEt2uU1g4qM"
      },
      "source": [
        "## Network"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 151,
      "metadata": {
        "id": "Gc6q2sg2JYeg"
      },
      "outputs": [],
      "source": [
        "INPUT_DIM = x_data.shape[-1]\n",
        "OUTPUT_DIM = y_data.shape[-1]\n",
        "\n",
        "# Linear regression model\n",
        "model = torch.nn.Sequential(\n",
        "    torch.nn.Linear(INPUT_DIM, OUTPUT_DIM),\n",
        "    )    \n",
        " "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gGVx0r32B9YK"
      },
      "source": [
        "List model parameters:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 152,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6vkVkh6Bh5mv",
        "outputId": "d494d89d-b1f9-4a1e-938b-bace4397b8f1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "0.weight: torch.Size([1, 13])\n",
            "0.bias: torch.Size([1])\n"
          ]
        }
      ],
      "source": [
        "for k, p in model.named_parameters():\n",
        "    print(f\"{k}: {p.shape}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Wrapper"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 153,
      "metadata": {},
      "outputs": [],
      "source": [
        "bnn = bnn_wrapper.BayesianNeuralNetwork(model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5yi11c7oV4x8"
      },
      "source": [
        "## Priors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 154,
      "metadata": {},
      "outputs": [],
      "source": [
        "bnn.set_prior_densities(densities.create_gaussian_nll)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Likelihood"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 155,
      "metadata": {},
      "outputs": [],
      "source": [
        "LIKELIHOOD_SIGMA = 100\n",
        "\n",
        "def predictive_distribution_sampler(logits, **kwargs): \n",
        "    return likelihoods.factorized_gaussian_with_fixed_scale_sample(logits, scale=LIKELIHOOD_SIGMA, **kwargs)\n",
        "\n",
        "def predictive_distribution_log_lik(logits, output_y, **kwargs): \n",
        "    return likelihoods.factorized_gaussian_with_fixed_scale_log_prob(logits, output_y, scale=LIKELIHOOD_SIGMA, **kwargs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 156,
      "metadata": {},
      "outputs": [],
      "source": [
        "bnn.predictive_distribution_sampler = predictive_distribution_sampler\n",
        "bnn.predictive_distribution_log_lik = predictive_distribution_log_lik"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9D7EgkWdruSn"
      },
      "source": [
        "## Create samplers for parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 157,
      "metadata": {},
      "outputs": [],
      "source": [
        "# # factorized gaussian\n",
        "# sampler, variational_params1, aux_objs = sampling.create_multiparameter_sampler(sampling.create_factorized_gaussian_sampler, model.named_parameters())\n",
        "\n",
        "# normalizing flow\n",
        "from reparameterized.sampling.realnvp import build_realnvp\n",
        "sampler, variational_params1, aux_objs = sampling.create_multiparameter_sampler(\n",
        "    sampling.create_flow_sampler, \n",
        "    model.named_parameters(),\n",
        "    build_flow_func=build_realnvp,\n",
        "    realnvp_m=8,  \n",
        "    realnvp_num_layers=16,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Note, we use only one sampler returning joint samples for all parameters.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 158,
      "metadata": {},
      "outputs": [],
      "source": [
        "parameter_names = tuple(n for n, _ in model.named_parameters())\n",
        "\n",
        "bnn.set_posterior_sampler(parameter_names, sampler, variational_params1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3_7JDyMt1R2W"
      },
      "source": [
        "# Preview sample"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 159,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "samplers:\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "{('0.weight',\n",
              "  '0.bias'): <function reparameterized.sampling.multiparameter._multiparameter_sampler_unpack.<locals>.wrapped_sampler(n_samples=1)>}"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "print(\"samplers:\")\n",
        "display(bnn.parameters2sampler)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 160,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "variational parameters:\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "[('0.weight_0.bias:t.0.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.0.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.0.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.0.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.0.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.0.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.1.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.1.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.1.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.1.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.1.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.1.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.2.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.2.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.2.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.2.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.2.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.2.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.3.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.3.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.3.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.3.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.3.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.3.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.4.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.4.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.4.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.4.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.4.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.4.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.5.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.5.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.5.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.5.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.5.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.5.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.6.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.6.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.6.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.6.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.6.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.6.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.7.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.7.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.7.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.7.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.7.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.7.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.8.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.8.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.8.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.8.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.8.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.8.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.9.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.9.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.9.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.9.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.9.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.9.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.10.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.10.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.10.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.10.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.10.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.10.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.11.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.11.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.11.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.11.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.11.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.11.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.12.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.12.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.12.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.12.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.12.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.12.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.13.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.13.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.13.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.13.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.13.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.13.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.14.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.14.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.14.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.14.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.14.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.14.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:t.15.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:t.15.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.15.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:t.15.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:t.15.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:t.15.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.0.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.0.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.0.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.0.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.0.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.0.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.1.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.1.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.1.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.1.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.1.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.1.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.2.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.2.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.2.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.2.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.2.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.2.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.3.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.3.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.3.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.3.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.3.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.3.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.4.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.4.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.4.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.4.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.4.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.4.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.5.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.5.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.5.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.5.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.5.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.5.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.6.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.6.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.6.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.6.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.6.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.6.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.7.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.7.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.7.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.7.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.7.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.7.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.8.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.8.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.8.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.8.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.8.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.8.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.9.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.9.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.9.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.9.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.9.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.9.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.10.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.10.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.10.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.10.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.10.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.10.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.11.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.11.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.11.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.11.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.11.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.11.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.12.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.12.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.12.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.12.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.12.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.12.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.13.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.13.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.13.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.13.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.13.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.13.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.14.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.14.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.14.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.14.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.14.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.14.4.bias', torch.Size([7])),\n",
              " ('0.weight_0.bias:s.15.0.weight', torch.Size([8, 7])),\n",
              " ('0.weight_0.bias:s.15.0.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.15.2.weight', torch.Size([8, 8])),\n",
              " ('0.weight_0.bias:s.15.2.bias', torch.Size([8])),\n",
              " ('0.weight_0.bias:s.15.4.weight', torch.Size([7, 8])),\n",
              " ('0.weight_0.bias:s.15.4.bias', torch.Size([7]))]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "print(\"variational parameters:\")\n",
        "display([(vn, vp.shape) for vn, vp in bnn.variational_params])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 161,
      "metadata": {
        "id": "WVRjL_Vb3Cyc"
      },
      "outputs": [],
      "source": [
        "parameters_samples, nlls = bnn.sample_posterior(n_samples=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 162,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nAElvBuKsutz",
        "outputId": "9906dfbf-fb99-44eb-c0f1-bbeedda8266d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "parameters' sample:\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "{'0.weight': torch.Size([3, 1, 13]), '0.bias': torch.Size([3, 1])}"
            ]
          },
          "execution_count": 162,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "print(\"parameters' sample:\")\n",
        "{n: s.shape for n, s in parameters_samples.items()}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 163,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cBNkS9FM1V6g",
        "outputId": "a72458d3-c921-45db-fb8d-deea62b85f9f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "parameters' NLLs:\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "tensor([[14.1378, 15.3474, 16.2944]], grad_fn=<StackBackward0>)"
            ]
          },
          "execution_count": 163,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "print(\"parameters' NLLs:\")\n",
        "nlls"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 164,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QkAbLF_UVB3s",
        "outputId": "a47fe134-088d-4be0-8104-e3a7fd4382ae"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "torch.Size([3, 2, 506, 1])"
            ]
          },
          "execution_count": 164,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "bnn.sample_predictive(x, n_samples=2, n_predictive_samples=3, flatten_samples_dims=False).shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 165,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wMK1M-vcYxkv",
        "outputId": "291b25f9-3f8f-40a1-e040-edde416bbabf"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "torch.Size([1, 506])"
            ]
          },
          "execution_count": 165,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "bnn.predictive_likelihoods(x, y).shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VqLQg7NvguxA"
      },
      "source": [
        "## Learning"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 166,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HgRQ6dDMJYer",
        "outputId": "b88b19fd-06f6-41e0-9836-95702e68bd30"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch=0: elbo = -5609.82 (= -5597.69 [+] -29.14 [-] -17.00)\n",
            "epoch=1: elbo = -5269.90 (= -5257.65 [+] -29.09 [-] -16.85)\n",
            "epoch=2: elbo = -6665.23 (= -6655.85 [+] -25.66 [-] -16.28)\n",
            "epoch=3: elbo = -5580.00 (= -5570.40 [+] -25.55 [-] -15.95)\n",
            "epoch=4: elbo = -4508.60 (= -4500.56 [+] -25.75 [-] -17.71)\n",
            "epoch=5: elbo = -4784.37 (= -4776.29 [+] -25.63 [-] -17.55)\n",
            "epoch=6: elbo = -4007.08 (= -3995.75 [+] -28.47 [-] -17.14)\n",
            "epoch=7: elbo = -4406.60 (= -4392.26 [+] -30.57 [-] -16.23)\n",
            "epoch=8: elbo = -5151.36 (= -5141.00 [+] -28.38 [-] -18.02)\n",
            "epoch=9: elbo = -4350.73 (= -4337.38 [+] -31.12 [-] -17.77)\n",
            "epoch=100: elbo = -2876.93 (= -2858.56 [+] -27.04 [-] -8.68)\n",
            "epoch=200: elbo = -2859.32 (= -2839.10 [+] -26.13 [-] -5.91)\n",
            "epoch=300: elbo = -2849.90 (= -2831.71 [+] -23.00 [-] -4.81)\n",
            "epoch=400: elbo = -2828.23 (= -2806.60 [+] -21.69 [-] -0.06)\n",
            "epoch=500: elbo = -2835.41 (= -2819.28 [+] -18.00 [-] -1.88)\n",
            "epoch=600: elbo = -2820.89 (= -2806.14 [+] -18.08 [-] -3.33)\n",
            "epoch=700: elbo = -2826.61 (= -2810.81 [+] -17.88 [-] -2.08)\n",
            "epoch=800: elbo = -2831.15 (= -2814.72 [+] -17.99 [-] -1.56)\n",
            "epoch=900: elbo = -2818.32 (= -2801.81 [+] -16.45 [-]  0.06)\n",
            "epoch=1000: elbo = -2819.91 (= -2802.99 [+] -16.38 [-]  0.55)\n",
            "epoch=1100: elbo = -2817.41 (= -2800.75 [+] -16.12 [-]  0.54)\n",
            "epoch=1200: elbo = -2820.42 (= -2803.30 [+] -16.24 [-]  0.88)\n",
            "epoch=1300: elbo = -2816.56 (= -2802.53 [+] -15.60 [-] -1.57)\n",
            "epoch=1400: elbo = -2816.81 (= -2801.79 [+] -15.58 [-] -0.56)\n",
            "epoch=1500: elbo = -2814.57 (= -2800.32 [+] -15.94 [-] -1.69)\n",
            "epoch=1600: elbo = -2817.00 (= -2801.35 [+] -15.66 [-] -0.01)\n",
            "epoch=1700: elbo = -2815.50 (= -2801.10 [+] -15.90 [-] -1.49)\n",
            "epoch=1800: elbo = -2814.79 (= -2801.49 [+] -16.17 [-] -2.87)\n",
            "epoch=1900: elbo = -2814.85 (= -2801.02 [+] -15.62 [-] -1.79)\n",
            "epoch=2000: elbo = -2813.59 (= -2798.46 [+] -15.67 [-] -0.54)\n",
            "epoch=2100: elbo = -2814.18 (= -2800.16 [+] -16.81 [-] -2.79)\n",
            "epoch=2200: elbo = -2814.70 (= -2800.90 [+] -16.36 [-] -2.56)\n",
            "epoch=2300: elbo = -2813.84 (= -2800.33 [+] -15.65 [-] -2.14)\n",
            "epoch=2400: elbo = -2813.60 (= -2801.06 [+] -15.08 [-] -2.54)\n",
            "epoch=2500: elbo = -2812.40 (= -2799.18 [+] -15.91 [-] -2.69)\n",
            "epoch=2600: elbo = -2813.20 (= -2800.18 [+] -15.44 [-] -2.42)\n",
            "epoch=2700: elbo = -2813.87 (= -2801.64 [+] -15.35 [-] -3.12)\n",
            "epoch=2800: elbo = -2813.46 (= -2799.68 [+] -15.92 [-] -2.15)\n",
            "epoch=2900: elbo = -2812.64 (= -2801.16 [+] -16.09 [-] -4.60)\n",
            "epoch=3000: elbo = -2812.93 (= -2800.13 [+] -16.27 [-] -3.48)\n",
            "epoch=3100: elbo = -2812.82 (= -2799.85 [+] -16.10 [-] -3.13)\n",
            "epoch=3200: elbo = -2812.25 (= -2799.65 [+] -15.86 [-] -3.26)\n",
            "epoch=3300: elbo = -2812.29 (= -2799.68 [+] -16.91 [-] -4.30)\n",
            "epoch=3400: elbo = -2811.56 (= -2799.31 [+] -16.01 [-] -3.77)\n",
            "epoch=3500: elbo = -2812.20 (= -2800.53 [+] -17.45 [-] -5.78)\n",
            "epoch=3600: elbo = -2811.34 (= -2799.01 [+] -15.62 [-] -3.30)\n",
            "epoch=3700: elbo = -2810.76 (= -2799.39 [+] -16.07 [-] -4.70)\n",
            "epoch=3800: elbo = -2810.25 (= -2798.62 [+] -16.39 [-] -4.77)\n",
            "epoch=3900: elbo = -2811.85 (= -2800.30 [+] -16.15 [-] -4.60)\n",
            "epoch=4000: elbo = -2811.12 (= -2799.79 [+] -15.30 [-] -3.97)\n",
            "epoch=4100: elbo = -2810.89 (= -2800.28 [+] -16.75 [-] -6.14)\n",
            "epoch=4200: elbo = -2811.33 (= -2798.89 [+] -15.83 [-] -3.39)\n",
            "epoch=4300: elbo = -2811.46 (= -2799.44 [+] -17.07 [-] -5.05)\n",
            "epoch=4400: elbo = -2811.26 (= -2799.58 [+] -16.28 [-] -4.60)\n",
            "epoch=4500: elbo = -2810.81 (= -2799.23 [+] -16.38 [-] -4.80)\n",
            "epoch=4600: elbo = -2811.00 (= -2799.17 [+] -16.16 [-] -4.33)\n",
            "epoch=4700: elbo = -2811.23 (= -2799.71 [+] -16.34 [-] -4.81)\n",
            "epoch=4800: elbo = -2811.40 (= -2800.15 [+] -16.13 [-] -4.89)\n",
            "epoch=4900: elbo = -2810.54 (= -2800.20 [+] -16.82 [-] -6.48)\n",
            "epoch=5000: elbo = -2810.85 (= -2799.41 [+] -16.73 [-] -5.30)\n",
            "epoch=5100: elbo = -2811.35 (= -2800.13 [+] -16.55 [-] -5.33)\n",
            "epoch=5200: elbo = -2810.80 (= -2800.58 [+] -17.47 [-] -7.24)\n",
            "epoch=5300: elbo = -2811.24 (= -2799.68 [+] -16.46 [-] -4.90)\n",
            "epoch=5400: elbo = -2811.16 (= -2799.88 [+] -16.01 [-] -4.73)\n",
            "epoch=5500: elbo = -2810.71 (= -2800.39 [+] -16.67 [-] -6.35)\n",
            "epoch=5600: elbo = -2810.80 (= -2799.18 [+] -16.59 [-] -4.98)\n",
            "epoch=5700: elbo = -2810.78 (= -2799.82 [+] -16.26 [-] -5.30)\n",
            "epoch=5800: elbo = -2810.43 (= -2800.07 [+] -16.71 [-] -6.36)\n",
            "epoch=5900: elbo = -2810.46 (= -2799.78 [+] -16.32 [-] -5.64)\n",
            "epoch=6000: elbo = -2810.61 (= -2800.40 [+] -17.36 [-] -7.15)\n",
            "epoch=6100: elbo = -2810.67 (= -2799.99 [+] -16.52 [-] -5.84)\n",
            "epoch=6200: elbo = -2810.11 (= -2799.23 [+] -17.01 [-] -6.13)\n",
            "epoch=6300: elbo = -2810.83 (= -2799.85 [+] -16.74 [-] -5.76)\n",
            "epoch=6400: elbo = -2809.76 (= -2799.19 [+] -15.86 [-] -5.29)\n",
            "epoch=6500: elbo = -2810.14 (= -2799.48 [+] -16.53 [-] -5.87)\n",
            "epoch=6600: elbo = -2810.78 (= -2799.67 [+] -17.57 [-] -6.45)\n",
            "epoch=6700: elbo = -2810.31 (= -2800.93 [+] -17.31 [-] -7.93)\n",
            "epoch=6800: elbo = -2810.48 (= -2799.86 [+] -16.17 [-] -5.55)\n",
            "epoch=6900: elbo = -2810.87 (= -2800.09 [+] -16.16 [-] -5.37)\n",
            "epoch=7000: elbo = -2810.74 (= -2800.16 [+] -16.34 [-] -5.76)\n",
            "epoch=7100: elbo = -2810.91 (= -2801.08 [+] -16.54 [-] -6.71)\n",
            "epoch=7200: elbo = -2810.35 (= -2799.76 [+] -16.77 [-] -6.18)\n",
            "epoch=7300: elbo = -2810.43 (= -2800.26 [+] -16.49 [-] -6.32)\n",
            "epoch=7400: elbo = -2810.33 (= -2798.87 [+] -16.65 [-] -5.19)\n",
            "epoch=7500: elbo = -2810.57 (= -2801.22 [+] -16.49 [-] -7.14)\n",
            "epoch=7600: elbo = -2810.54 (= -2800.36 [+] -17.31 [-] -7.13)\n",
            "epoch=7700: elbo = -2810.76 (= -2799.79 [+] -16.38 [-] -5.41)\n",
            "epoch=7800: elbo = -2810.26 (= -2800.01 [+] -16.52 [-] -6.27)\n",
            "epoch=7900: elbo = -2811.37 (= -2800.35 [+] -17.11 [-] -6.09)\n",
            "epoch=8000: elbo = -2811.80 (= -2800.99 [+] -17.15 [-] -6.34)\n",
            "epoch=8100: elbo = -2810.36 (= -2799.58 [+] -16.23 [-] -5.45)\n",
            "epoch=8200: elbo = -2810.58 (= -2799.56 [+] -16.71 [-] -5.70)\n",
            "epoch=8300: elbo = -2810.82 (= -2799.35 [+] -16.44 [-] -4.98)\n",
            "epoch=8400: elbo = -2810.62 (= -2799.52 [+] -16.36 [-] -5.27)\n",
            "epoch=8500: elbo = -2810.68 (= -2799.33 [+] -17.22 [-] -5.88)\n",
            "epoch=8600: elbo = -2810.77 (= -2799.50 [+] -16.40 [-] -5.13)\n",
            "epoch=8700: elbo = -2810.68 (= -2800.13 [+] -16.50 [-] -5.95)\n",
            "epoch=8800: elbo = -2810.67 (= -2799.00 [+] -16.90 [-] -5.23)\n",
            "epoch=8900: elbo = -2810.36 (= -2800.88 [+] -16.38 [-] -6.90)\n",
            "epoch=9000: elbo = -2810.43 (= -2799.80 [+] -16.70 [-] -6.07)\n",
            "epoch=9100: elbo = -2810.72 (= -2800.25 [+] -17.82 [-] -7.35)\n",
            "epoch=9200: elbo = -2810.33 (= -2799.30 [+] -16.35 [-] -5.32)\n",
            "epoch=9300: elbo = -2810.92 (= -2800.57 [+] -17.20 [-] -6.85)\n",
            "epoch=9400: elbo = -2810.40 (= -2799.82 [+] -15.82 [-] -5.25)\n",
            "epoch=9500: elbo = -2810.44 (= -2800.22 [+] -16.60 [-] -6.39)\n",
            "epoch=9600: elbo = -2810.66 (= -2800.26 [+] -16.33 [-] -5.93)\n",
            "epoch=9700: elbo = -2810.20 (= -2799.08 [+] -16.27 [-] -5.15)\n",
            "epoch=9800: elbo = -2810.05 (= -2799.90 [+] -16.73 [-] -6.57)\n",
            "epoch=9900: elbo = -2810.14 (= -2800.61 [+] -16.19 [-] -6.66)\n"
          ]
        }
      ],
      "source": [
        "optimized_parameters = [vp for _, vp in bnn.variational_params]\n",
        "optimizer = torch.optim.Adam(optimized_parameters, lr=0.001) \n",
        "\n",
        "n_posterior_samples = 10\n",
        "n_epochs = 10000\n",
        "\n",
        "for e in range(n_epochs):\n",
        "    optimizer.zero_grad()\n",
        "\n",
        "    # sampling from approximate posterior\n",
        "    parameters_samples, posterior_nlls = bnn.sample_posterior(n_samples=n_posterior_samples)\n",
        "    log_q = -posterior_nlls.sum(0)  # sum over parameters\n",
        "\n",
        "    # priors\n",
        "    prior_nlls = bnn.prior_nll(parameters_samples)\n",
        "    log_p = -prior_nlls.sum(0)  # sum over parameters\n",
        "\n",
        "    # estimating likelihoods of observed y\n",
        "    logliks = bnn.predictive_likelihoods(x, y, parameters_samples)\n",
        "    assert logliks.shape==torch.Size([n_posterior_samples, len(y)])\n",
        "    logliks = logliks.sum(1)  # sum over data points\n",
        "    \n",
        "    elbo = logliks+log_p-log_q\n",
        "    elbo = elbo.mean()  # average over posterior samples\n",
        "\n",
        "    loss_vi = -elbo\n",
        "    loss_vi.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    if e<10 or e%100==0:\n",
        "        print(f\"epoch={e}: elbo = {elbo: .2f} (= {logliks.mean(): .2f} [+] {log_p.mean(): .2f} [-] {log_q.mean(): .2f})\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FMLYtaZFBoFm"
      },
      "source": [
        "# Linear Regression Bayesian true (closed-form) posterior"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O_qCizafZv1z"
      },
      "source": [
        "See Bishop's book"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 167,
      "metadata": {
        "id": "1yIXrPEr7dDU"
      },
      "outputs": [],
      "source": [
        "psi = np.hstack([x_data, np.ones(x_data.shape)])\n",
        "t = y_data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 168,
      "metadata": {
        "id": "elCPvgtiEQwf"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "((506, 26), (506, 1))"
            ]
          },
          "execution_count": 168,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "psi.shape, t.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 169,
      "metadata": {
        "id": "KH-ZfaL-_hwr"
      },
      "outputs": [],
      "source": [
        "#likelihood precision:\n",
        "beta = 1./(LIKELIHOOD_SIGMA*LIKELIHOOD_SIGMA)\n",
        "\n",
        "#prior:\n",
        "S0 = np.eye(psi.shape[-1])\n",
        "m0 = np.zeros(psi.shape[-1]).reshape(-1, 1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 170,
      "metadata": {
        "id": "wsl3q5Ys7_sJ"
      },
      "outputs": [],
      "source": [
        "SN = np.linalg.inv( np.linalg.inv(S0) + beta*np.dot(psi.T, psi) ) # covariance matrix (Bishop 3.50)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 171,
      "metadata": {
        "id": "YPZEyITg_uld"
      },
      "outputs": [],
      "source": [
        "mN = np.dot(SN, np.dot(np.linalg.inv(S0), m0) + beta*np.dot(psi.T, t)) # means (Bishop 3.51)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compare expectations (means):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 172,
      "metadata": {
        "id": "jqLCbzz5AmTr"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "array([-0.09,  0.12, -0.06,  0.02,  0.01,  0.28,  0.09,  0.02,  0.03,\n",
              "        0.  ,  0.18,  0.04, -0.47,  0.02])"
            ]
          },
          "execution_count": 172,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dim = SN.shape[0]//2+1\n",
        "\n",
        "mN[:dim].flatten().round(2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 173,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(tensor([-0.1200,  0.0900, -0.0600,  0.0100,  0.1500,  0.3400,  0.0600,  0.2300,\n",
              "          0.0600,  0.0000,  0.1000,  0.0500, -0.4800], grad_fn=<RoundBackward1>),\n",
              " tensor(-0.0092, grad_fn=<MeanBackward0>))"
            ]
          },
          "execution_count": 173,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# empirically estimate means for the VI solution:\n",
        "parameters_samples, nlls = bnn.sample_posterior(n_samples=10000)\n",
        "torch.round(parameters_samples['0.weight'].mean(dim=0).flatten(), decimals=2), parameters_samples['0.bias'].mean()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compare covariance:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 174,
      "metadata": {
        "id": "pj-hoGZ1A9Ba"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "array([[ 0.306, -0.003,  0.015,  0.001,  0.   ,  0.005, -0.003,  0.013,\n",
              "        -0.065, -0.004,  0.005,  0.004, -0.04 ,  0.001],\n",
              "       [-0.003,  0.051,  0.036, -0.   , -0.   , -0.018,  0.014, -0.037,\n",
              "         0.021, -0.002, -0.001, -0.002,  0.005, -0.002],\n",
              "       [ 0.015,  0.036,  0.564, -0.003, -0.002,  0.012, -0.027,  0.04 ,\n",
              "         0.062, -0.014,  0.008,  0.002, -0.044,  0.   ],\n",
              "       [ 0.001, -0.   , -0.003,  0.997, -0.   , -0.001, -0.001,  0.001,\n",
              "        -0.002,  0.   ,  0.002, -0.   ,  0.003, -0.   ],\n",
              "       [ 0.   , -0.   , -0.002, -0.   ,  1.   , -0.002, -0.002, -0.   ,\n",
              "         0.001, -0.   , -0.003, -0.   , -0.001, -0.   ],\n",
              "       [ 0.005, -0.018,  0.012, -0.001, -0.002,  0.946, -0.018, -0.026,\n",
              "         0.014, -0.003, -0.07 , -0.007,  0.028, -0.005],\n",
              "       [-0.003,  0.014, -0.027, -0.001, -0.002, -0.018,  0.044,  0.018,\n",
              "         0.01 , -0.003, -0.018, -0.003, -0.041, -0.002],\n",
              "       [ 0.013, -0.037,  0.04 ,  0.001, -0.   , -0.026,  0.018,  0.908,\n",
              "         0.022, -0.002, -0.083, -0.006, -0.019, -0.004],\n",
              "       [-0.065,  0.021,  0.062, -0.002,  0.001,  0.014,  0.01 ,  0.022,\n",
              "         0.615, -0.027,  0.029,  0.007,  0.02 ,  0.003],\n",
              "       [-0.004, -0.002, -0.014,  0.   , -0.   , -0.003, -0.003, -0.002,\n",
              "        -0.027,  0.003, -0.011, -0.   , -0.004, -0.001],\n",
              "       [ 0.005, -0.001,  0.008,  0.002, -0.003, -0.07 , -0.018, -0.083,\n",
              "         0.029, -0.011,  0.697, -0.016, -0.05 , -0.011],\n",
              "       [ 0.004, -0.002,  0.002, -0.   , -0.   , -0.007, -0.003, -0.006,\n",
              "         0.007, -0.   , -0.016,  0.002,  0.001, -0.001],\n",
              "       [-0.04 ,  0.005, -0.044,  0.003, -0.001,  0.028, -0.041, -0.019,\n",
              "         0.02 , -0.004, -0.05 ,  0.001,  0.424, -0.002],\n",
              "       [ 0.001, -0.002,  0.   , -0.   , -0.   , -0.005, -0.002, -0.004,\n",
              "         0.003, -0.001, -0.011, -0.001, -0.002,  0.999]])"
            ]
          },
          "execution_count": 174,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dim = SN.shape[0]//2+1\n",
        "\n",
        "SN = SN[:dim, :dim]\n",
        "\n",
        "SN.round(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 175,
      "metadata": {},
      "outputs": [],
      "source": [
        "# empirically estimate covariance matrix for the VI solution:\n",
        "parameters_samples, nlls = bnn.sample_posterior(n_samples=10)\n",
        "samples_weight, samples_bias = parameters_samples['0.weight'].flatten(start_dim=1), parameters_samples['0.bias'].flatten(start_dim=1)\n",
        "samples = torch.hstack([samples_weight, samples_bias]).T\n",
        "cov = torch.cov(samples)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 176,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(array([0.306, 0.051, 0.564, 0.997, 1.   , 0.946, 0.044, 0.908, 0.615,\n",
              "        0.003, 0.696, 0.002, 0.424, 0.999]),\n",
              " array([6.900e-02, 5.100e-02, 9.380e-01, 2.240e-01, 1.575e+00, 1.535e+00,\n",
              "        5.400e-02, 9.260e-01, 3.250e-01, 1.000e-03, 3.250e-01, 1.000e-03,\n",
              "        3.360e-01, 1.385e+00]))"
            ]
          },
          "execution_count": 176,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Variances\n",
        "var1 = np.array([SN[i,i].round(4) for i in range(SN.shape[0])])\n",
        "var2 = np.array([float(cov[i,i]) for i in range(SN.shape[0])])\n",
        "\n",
        "var1.round(3), var2.round(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "-C2DLYqynn54"
      ],
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3.8.5 ('base')",
      "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.5"
    },
    "vscode": {
      "interpreter": {
        "hash": "72dff520818572b45c401a256658ee25fa99b55260abeb0336ef1b1e43b37267"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}