working locally

.gitignore

@@ -0,0 +1,3 @@
+.ipynb_checkpoints
+flagged
+.idea

SettingUpTheWebiste.ipynb

@@ -0,0 +1,209 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "963a04b2",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8ebc97aa",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "b73f00ce",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Running locally at: http://127.0.0.1:7860/\n",
+      "To create a public link, set `share=True` in `launch()`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "\n",
+       "        <iframe\n",
+       "            width=\"900\"\n",
+       "            height=\"500\"\n",
+       "            src=\"http://127.0.0.1:7860/\"\n",
+       "            frameborder=\"0\"\n",
+       "            allowfullscreen\n",
+       "        ></iframe>\n",
+       "        "
+      ],
+      "text/plain": [
+       "<IPython.lib.display.IFrame at 0x7f8f67cba520>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "(<Flask 'gradio.networking'>, 'http://127.0.0.1:7860/', None)"
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import gradio as gr\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "import gpytorch\n",
+    "import torch\n",
+    "import sys\n",
+    "\n",
+    "import gpytorch\n",
+    "\n",
+    "# We will use the simplest form of GP model, exact inference\n",
+    "class ExactGPModel(gpytorch.models.ExactGP):\n",
+    "    def __init__(self, train_x, train_y, likelihood):\n",
+    "        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)\n",
+    "        self.mean_module = gpytorch.means.ConstantMean()\n",
+    "        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        mean_x = self.mean_module(x)\n",
+    "        covar_x = self.covar_module(x)\n",
+    "        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)\n",
+    "\n",
+    "def get_model(x, y, hyperparameters):\n",
+    "    likelihood = gpytorch.likelihoods.GaussianLikelihood(noise_constraint=gpytorch.constraints.GreaterThan(1.e-9))\n",
+    "    model = ExactGPModel(x, y, likelihood)\n",
+    "    model.likelihood.noise = torch.ones_like(model.likelihood.noise) * hyperparameters[\"noise\"]\n",
+    "    model.covar_module.outputscale = torch.ones_like(model.covar_module.outputscale) * hyperparameters[\"outputscale\"]\n",
+    "    model.covar_module.base_kernel.lengthscale = torch.ones_like(model.covar_module.base_kernel.lengthscale) * \\\n",
+    "                                                 hyperparameters[\"lengthscale\"]\n",
+    "    return model, likelihood\n",
+    "\n",
+    "\n",
+    "\n",
+    "excuse = \"Please only specify numbers, x values should be in [0,1] and y values in [-1,1].\"\n",
+    "excuse_max_examples = \"This model is trained to work with up to 4 input points.\"\n",
+    "hyperparameters = {'noise': 1e-4, 'outputscale': 1., 'lengthscale': .1, 'fast_computations': (False,False,False)}\n",
+    "\n",
+    "\n",
+    "conf = .5\n",
+    "\n",
+    "def mean_and_bounds_for_gp(x,y,test_xs):\n",
+    "    gp_model, likelihood = get_model(x,y,hyperparameters)\n",
+    "    gp_model.eval()\n",
+    "    l = likelihood(gp_model(test_xs))\n",
+    "    means = l.mean.squeeze()\n",
+    "    varis = torch.diagonal(l.covariance_matrix.squeeze())\n",
+    "    stds = varis.sqrt()\n",
+    "    return means, means-stds, means+stds\n",
+    "\n",
+    "\n",
+    "def mean_and_bounds_for_pnf(x,y,test_xs, choice):\n",
+    "    sys.path.append('prior-fitting/')\n",
+    "    model = torch.load(f'onefeature_gp_ls.1_pnf_{choice}.pt')\n",
+    "\n",
+    "\n",
+    "    logits = model((torch.cat([x,test_xs],0).unsqueeze(1),y.unsqueeze(1)),single_eval_pos=len(x))\n",
+    "    bounds = model.criterion.quantile(logits,center_prob=.682).squeeze(1)\n",
+    "    return model.criterion.mean(logits).squeeze(1), bounds[:,0], bounds[:,1]\n",
+    "\n",
+    "def plot_w_conf_interval(ax_or_plt, x, m, lb, ub, color):\n",
+    "    ax_or_plt.plot(x.squeeze(-1),m, color=color)\n",
+    "    ax_or_plt.fill_between(x.squeeze(-1), lb, ub, alpha=.1, color=color)\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "@torch.no_grad()\n",
+    "def infer(table, choice):\n",
+    "    vfunc = np.vectorize(lambda s: len(s))\n",
+    "    non_empty_row_mask = (vfunc(table).sum(1) != 0)\n",
+    "    table = table[non_empty_row_mask]\n",
+    "\n",
+    "    try:\n",
+    "        table = table.astype(np.float32)\n",
+    "    except ValueError:\n",
+    "        return excuse, None\n",
+    "    x = torch.tensor(table[:,0]).unsqueeze(1)\n",
+    "    y = torch.tensor(table[:,1])\n",
+    "    fig = plt.figure()\n",
+    "\n",
+    "    if len(x) > 4:\n",
+    "        return excuse_max_examples, None\n",
+    "    if (x<0.).any() or (x>1.).any() or (y<-1).any() or (y>1).any():\n",
+    "        return excuse, None\n",
+    "\n",
+    "    plt.scatter(x,y)\n",
+    "\n",
+    "\n",
+    "    \n",
+    "    test_xs = torch.linspace(0,1,100).unsqueeze(1)\n",
+    "    \n",
+    "    plot_w_conf_interval(plt, test_xs, *mean_and_bounds_for_gp(x,y,test_xs), 'green')\n",
+    "    plot_w_conf_interval(plt, test_xs, *mean_and_bounds_for_pnf(x,y,test_xs, choice), 'blue')\n",
+    "\n",
+    "\n",
+    "    \n",
+    "    return '', plt.gcf()\n",
+    "\n",
+    "iface = gr.Interface(fn=infer, \n",
+    "                     inputs=[\n",
+    "                         gr.inputs.Dataframe(headers=[\"x\", \"y\"], datatype=[\"number\", \"number\"], row_count=2, type='numpy', default=[['.25','.1'],['.75','.4']]),\n",
+    "                         gr.inputs.Radio(['160K','800K','4M'], type=\"value\", default='4M', label='Training Costs')\n",
+    "                     ], outputs=[\"text\",\"plot\"])\n",
+    "iface.launch()\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a3a377e3",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "72c0c821",
+   "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.9.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

app.py

@@ -1,7 +1,106 @@
 import gradio as gr
+import numpy as np
+import matplotlib.pyplot as plt
+import gpytorch
+import torch
+import sys
 
-def greet(name):
-    return "Hello " + name + "!!"
+import gpytorch
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+# We will use the simplest form of GP model, exact inference
+class ExactGPModel(gpytorch.models.ExactGP):
+    def __init__(self, train_x, train_y, likelihood):
+        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
+        self.mean_module = gpytorch.means.ConstantMean()
+        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
+
+    def forward(self, x):
+        mean_x = self.mean_module(x)
+        covar_x = self.covar_module(x)
+        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
+
+def get_model(x, y, hyperparameters):
+    likelihood = gpytorch.likelihoods.GaussianLikelihood(noise_constraint=gpytorch.constraints.GreaterThan(1.e-9))
+    model = ExactGPModel(x, y, likelihood)
+    model.likelihood.noise = torch.ones_like(model.likelihood.noise) * hyperparameters["noise"]
+    model.covar_module.outputscale = torch.ones_like(model.covar_module.outputscale) * hyperparameters["outputscale"]
+    model.covar_module.base_kernel.lengthscale = torch.ones_like(model.covar_module.base_kernel.lengthscale) * \
+                                                 hyperparameters["lengthscale"]
+    return model, likelihood
+
+
+
+excuse = "Please only specify numbers, x values should be in [0,1] and y values in [-1,1]."
+excuse_max_examples = "This model is trained to work with up to 4 input points."
+hyperparameters = {'noise': 1e-4, 'outputscale': 1., 'lengthscale': .1, 'fast_computations': (False,False,False)}
+
+
+conf = .5
+
+def mean_and_bounds_for_gp(x,y,test_xs):
+    gp_model, likelihood = get_model(x,y,hyperparameters)
+    gp_model.eval()
+    l = likelihood(gp_model(test_xs))
+    means = l.mean.squeeze()
+    varis = torch.diagonal(l.covariance_matrix.squeeze())
+    stds = varis.sqrt()
+    return means, means-stds, means+stds
+
+
+def mean_and_bounds_for_pnf(x,y,test_xs, choice):
+    sys.path.append('prior-fitting/')
+    model = torch.load(f'onefeature_gp_ls.1_pnf_{choice}.pt')
+
+
+    logits = model((torch.cat([x,test_xs],0).unsqueeze(1),y.unsqueeze(1)),single_eval_pos=len(x))
+    bounds = model.criterion.quantile(logits,center_prob=.682).squeeze(1)
+    return model.criterion.mean(logits).squeeze(1), bounds[:,0], bounds[:,1]
+
+def plot_w_conf_interval(ax_or_plt, x, m, lb, ub, color):
+    ax_or_plt.plot(x.squeeze(-1),m, color=color)
+    ax_or_plt.fill_between(x.squeeze(-1), lb, ub, alpha=.1, color=color)
+
+
+
+
+@torch.no_grad()
+def infer(table, choice):
+    vfunc = np.vectorize(lambda s: len(s))
+    non_empty_row_mask = (vfunc(table).sum(1) != 0)
+    table = table[non_empty_row_mask]
+
+    try:
+        table = table.astype(np.float32)
+    except ValueError:
+        return excuse, None
+    x = torch.tensor(table[:,0]).unsqueeze(1)
+    y = torch.tensor(table[:,1])
+    fig = plt.figure()
+
+    if len(x) > 4:
+        return excuse_max_examples, None
+    if (x<0.).any() or (x>1.).any() or (y<-1).any() or (y>1).any():
+        return excuse, None
+
+    plt.scatter(x,y)
+
+
+    
+    test_xs = torch.linspace(0,1,100).unsqueeze(1)
+    
+    plot_w_conf_interval(plt, test_xs, *mean_and_bounds_for_gp(x,y,test_xs), 'green')
+    plot_w_conf_interval(plt, test_xs, *mean_and_bounds_for_pnf(x,y,test_xs, choice), 'blue')
+
+
+    
+    return '', plt.gcf()
+
+iface = gr.Interface(fn=infer, 
+                     inputs=[
+                         gr.inputs.Dataframe(headers=["x", "y"], datatype=["number", "number"], row_count=2, type='numpy', default=[['.25','.1'],['.75','.4']]),
+                         gr.inputs.Radio(['160K','800K','4M'], type="value", default='4M', label='Training Costs')
+                     ], outputs=["text","plot"])
 iface.launch()
+
+
+

prior-fitting/.gitignore

@@ -0,0 +1,129 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

prior-fitting/README.md

@@ -0,0 +1 @@
+# prior-fitting

prior-fitting/acquisition_functions.py

@@ -0,0 +1,18 @@
+from botorch.acquisition import AcquisitionFunction
+from torch import Tensor
+
+
+class ExpectedImprovement(AcquisitionFunction):
+    def forward(self, X: Tensor, best_f, maximize=True) -> Tensor: # X: evaluation_points x feature_dim
+        assert len(X.shape) == 2
+
+        model = self.get_submodule('model')
+
+        y = model(X)
+
+        full_range = model.full_range
+
+
+
+
+

prior-fitting/bar_distribution.py

@@ -0,0 +1,147 @@
+
+import torch
+from torch import nn
+
+class BarDistribution(nn.Module):
+    def __init__(self, borders: torch.Tensor): # here borders should start with min and end with max, where all values lie in (min,max) and are sorted
+        # sorted list of borders
+        super().__init__()
+        assert len(borders.shape) == 1
+        #self.borders = borders
+        self.register_buffer('borders', borders)
+        #self.bucket_widths = self.borders[1:] - self.borders[:-1]
+        self.register_buffer('bucket_widths', self.borders[1:] - self.borders[:-1])
+        full_width = self.bucket_widths.sum()
+        assert (full_width - (self.borders[-1] - self.borders[0])).abs() < 1e-4, f'diff: {full_width - (self.borders[-1] - self.borders[0])}'
+        assert (torch.argsort(borders) == torch.arange(len(borders))).all(), "Please provide sorted borders!"
+        self.num_bars = len(borders) - 1
+
+    def map_to_bucket_idx(self, y):
+        target_sample = torch.searchsorted(self.borders, y) - 1
+        target_sample[y == self.borders[0]] = 0
+        target_sample[y == self.borders[-1]] = self.num_bars - 1
+        return target_sample
+
+    def forward(self, logits, y): # gives the negative log density (the _loss_), y: T x B, logits: T x B x self.num_bars
+        target_sample = self.map_to_bucket_idx(y)
+        assert (target_sample >= 0).all() and (target_sample < self.num_bars).all(), f'y {y} not in support set for borders (min_y, max_y) {self.borders}'
+        assert logits.shape[-1] == self.num_bars, f'{logits.shape[-1]} vs {self.num_bars}'
+
+        bucket_log_probs = torch.log_softmax(logits, -1)
+        scaled_bucket_log_probs = bucket_log_probs - torch.log(self.bucket_widths)
+
+        return -scaled_bucket_log_probs.gather(-1,target_sample.unsqueeze(-1)).squeeze(-1)
+
+    def mean(self, logits):
+        bucket_means = self.borders[:-1] + self.bucket_widths/2
+        p = torch.softmax(logits, -1)
+        return p @ bucket_means
+
+    def quantile(self, logits, center_prob=.682):
+        logits_shape = logits.shape
+        logits = logits.view(-1, logits.shape[-1])
+        side_prob = (1-center_prob)/2
+        probs = logits.softmax(-1)
+        flipped_probs = probs.flip(-1)
+        cumprobs = torch.cumsum(probs, -1)
+        flipped_cumprobs = torch.cumsum(flipped_probs, -1)
+
+        def find_lower_quantile(probs, cumprobs, side_prob, borders):
+            idx = (torch.searchsorted(cumprobs, side_prob)).clamp(0, len(cumprobs)-1) # this might not do the right for outliers
+
+            left_prob = cumprobs[idx-1]
+            rest_prob = side_prob - left_prob
+            left_border, right_border = borders[idx:idx+2]
+            return left_border + (right_border-left_border)*rest_prob/probs[idx]
+
+        results = []
+        for p,cp,f_p,f_cp in zip(probs, cumprobs, flipped_probs, flipped_cumprobs):
+            r = find_lower_quantile(p, cp, side_prob, self.borders), find_lower_quantile(f_p, f_cp, side_prob, self.borders.flip(0))
+            results.append(r)
+
+        return torch.tensor(results).reshape(*logits_shape[:-1],2)
+
+    def mode(self, logits):
+        mode_inds = logits.argmax(-1)
+        bucket_means = self.borders[:-1] + self.bucket_widths/2
+        return bucket_means[mode_inds]
+
+    def ei(self, logits, best_f, maximize=True): # logits: evaluation_points x batch x feature_dim
+        bucket_means = self.borders[:-1] + self.bucket_widths/2
+        if maximize:
+            bucket_contributions = torch.tensor(
+                [max((bucket_max + max(bucket_min, best_f)) / 2 - best_f,0) for
+                 bucket_min, bucket_max, bucket_mean in zip(self.borders[:-1], self.borders[1:], bucket_means)], dtype=logits.dtype, device=logits.device)
+        else:
+            bucket_contributions = torch.tensor(
+                [-min((min(bucket_max,best_f) + bucket_min) / 2 - best_f,0) for # min on max instead of max on min, and compare min < instead of max >
+                 bucket_min, bucket_max, bucket_mean in zip(self.borders[:-1], self.borders[1:], bucket_means)], dtype=logits.dtype, device=logits.device)
+        p = torch.softmax(logits, -1)
+        return p @ bucket_contributions
+
+
+class FullSupportBarDistribution(BarDistribution):
+    @staticmethod
+    def halfnormal_with_p_weight_before(range_max,p=.5):
+        s = range_max / torch.distributions.HalfNormal(torch.tensor(1.)).icdf(torch.tensor(p))
+        return torch.distributions.HalfNormal(s)
+
+    def forward(self, logits, y): # gives the negative log density (the _loss_), y: T x B, logits: T x B x self.num_bars
+        assert self.num_bars > 1
+        target_sample = self.map_to_bucket_idx(y)
+        target_sample.clamp_(0,self.num_bars-1)
+        assert logits.shape[-1] == self.num_bars
+
+        bucket_log_probs = torch.log_softmax(logits, -1)
+        scaled_bucket_log_probs = bucket_log_probs - torch.log(self.bucket_widths)
+        #print(bucket_log_probs, logits.shape)
+        log_probs = scaled_bucket_log_probs.gather(-1,target_sample.unsqueeze(-1)).squeeze(-1)
+
+        side_normals = (self.halfnormal_with_p_weight_before(self.bucket_widths[0]), self.halfnormal_with_p_weight_before(self.bucket_widths[-1]))
+
+
+        # TODO look over it again
+        log_probs[target_sample == 0] += side_normals[0].log_prob((self.borders[1]-y[target_sample == 0]).clamp(min=.00000001)) + torch.log(self.bucket_widths[0])
+        log_probs[target_sample == self.num_bars-1] += side_normals[1].log_prob(y[target_sample == self.num_bars-1]-self.borders[-2]) + torch.log(self.bucket_widths[-1])
+
+
+        return -log_probs
+
+    def mean(self, logits):
+        bucket_means = self.borders[:-1] + self.bucket_widths / 2
+        p = torch.softmax(logits, -1)
+        side_normals = (self.halfnormal_with_p_weight_before(self.bucket_widths[0]),
+                        self.halfnormal_with_p_weight_before(self.bucket_widths[-1]))
+        bucket_means[0] = -side_normals[0].mean + self.borders[1]
+        bucket_means[-1] = side_normals[1].mean + self.borders[-2]
+        return p @ bucket_means
+
+
+
+def get_bucket_limits(num_outputs:int, full_range:tuple=None, ys:torch.Tensor=None):
+    assert (ys is not None) or (full_range is not None)
+    if ys is not None:
+        ys = ys.flatten()
+        if len(ys) % num_outputs: ys = ys[:-(len(ys) % num_outputs)]
+        print(f'Using {len(ys)} y evals to estimate {num_outputs} buckets. Cut off the last {len(ys) % num_outputs} ys.')
+        ys_per_bucket = len(ys) // num_outputs
+        if full_range is None:
+            full_range = (ys.min(), ys.max())
+        else:
+            assert full_range[0] <= ys.min() and full_range[1] >= ys.max()
+            full_range = torch.tensor(full_range)
+        ys_sorted, ys_order = ys.sort(0)
+        bucket_limits = (ys_sorted[ys_per_bucket-1::ys_per_bucket][:-1]+ys_sorted[ys_per_bucket::ys_per_bucket])/2
+        print(full_range)
+        bucket_limits = torch.cat([full_range[0].unsqueeze(0), bucket_limits, full_range[1].unsqueeze(0)],0)
+
+    else:
+        class_width = (full_range[1] - full_range[0]) / num_outputs
+        bucket_limits = torch.cat([full_range[0] + torch.arange(num_outputs).float()*class_width, torch.tensor(full_range[1]).unsqueeze(0)], 0)
+
+    assert len(bucket_limits) - 1 == num_outputs and full_range[0] == bucket_limits[0] and full_range[-1] == bucket_limits[-1]
+    return bucket_limits
+
+
+
+

prior-fitting/decoders.py

@@ -0,0 +1,30 @@
+import torch
+from torch import nn
+import random
+
+
+class ScaledDecoder(nn.Module):
+    def __init__(self, ninp, nhid, nout):
+        super().__init__()
+        self.linear = nn.Linear(ninp, nhid)
+        self.linear1 = nn.Linear(nhid, nout)
+        self.linear2 = nn.Linear(nhid, 10)
+
+    def forward(self, x):
+        #return torch.cat([self.linear1(x), self.linear2(x)], -1)
+        x = self.linear(x)
+        x = nn.GELU()(x)
+        temps = self.linear2(x).softmax(-1) @ torch.tensor([1.,1.4,1.7,2.,5.,10.,20.,40.,80.,160.], device=x.device)
+        if random.random() > .99:
+            print(temps.shape,temps[:,:2])
+        return self.linear1(x) / temps.unsqueeze(-1)
+
+class FixedScaledDecoder(nn.Module):
+    def __init__(self, ninp, nhid, nout):
+        super().__init__()
+        self.mapper = nn.Sequential(nn.Linear(ninp, nhid), nn.GELU(), nn.Linear(nhid, nout))
+        self.T = nn.Parameter(torch.ones(10000)/10000)
+
+    def forward(self, x):
+        return self.mapper(x)/self.T.sum()
+

prior-fitting/encoders.py

@@ -0,0 +1,95 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
+
+class _PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout=0.):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        self.d_model = d_model
+        self.device_test_tensor = nn.Parameter(torch.tensor(1.))
+
+    def forward(self, x):# T x B x num_features
+        assert self.d_model % x.shape[-1]*2 == 0
+        d_per_feature = self.d_model // x.shape[-1]
+        pe = torch.zeros(*x.shape, d_per_feature, device=self.device_test_tensor.device)
+        #position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        interval_size = 10
+        div_term = (1./interval_size) * 2*math.pi*torch.exp(torch.arange(0, d_per_feature, 2, device=self.device_test_tensor.device).float()*math.log(math.sqrt(2)))
+        #print(div_term/2/math.pi)
+        pe[..., 0::2] = torch.sin(x.unsqueeze(-1) * div_term)
+        pe[..., 1::2] = torch.cos(x.unsqueeze(-1) * div_term)
+        return self.dropout(pe).view(x.shape[0],x.shape[1],self.d_model)
+
+
+class EmbeddingEncoder(nn.Module):
+    def __init__(self, num_features, em_size, num_embs=100):
+        super().__init__()
+        self.num_embs = num_embs
+        self.embeddings = nn.Embedding(num_embs * num_features, em_size, max_norm=True)
+        self.init_weights(.1)
+        self.min_max = (-2,+2)
+
+    @property
+    def width(self):
+        return self.min_max[1] - self.min_max[0]
+
+    def init_weights(self, initrange):
+        self.embeddings.weight.data.uniform_(-initrange, initrange)
+
+    def discretize(self, x):
+        split_size = self.width / self.num_embs
+        return (x - self.min_max[0] // split_size).int().clamp(0, self.num_embs - 1)
+
+    def forward(self, x):  # T x B x num_features
+        x_idxs = self.discretize(x)
+        x_idxs += torch.arange(x.shape[-1], device=x.device).view(1, 1, -1) * self.num_embs
+        # print(x_idxs,self.embeddings.weight.shape)
+        return self.embeddings(x_idxs).mean(-2)
+
+Linear = nn.Linear
+MLP = lambda num_features, emsize: nn.Sequential(nn.Linear(num_features+1,emsize*2),
+                                                 nn.ReLU(),
+                                                 nn.Linear(emsize*2,emsize))
+
+class Conv(nn.Module):
+    def __init__(self, input_size, emsize):
+        super().__init__()
+        self.convs = torch.nn.ModuleList([nn.Conv2d(64 if i else 1, 64, 3) for i in range(5)])
+        self.linear = nn.Linear(64,emsize)
+
+
+    def forward(self, x):
+        size = math.isqrt(x.shape[-1])
+        assert size*size == x.shape[-1]
+        x = x.reshape(*x.shape[:-1], 1, size, size)
+        for conv in self.convs:
+            if x.shape[-1] < 4:
+                break
+            x = conv(x)
+            x.relu_()
+        x = nn.AdaptiveAvgPool2d((1,1))(x).squeeze(-1).squeeze(-1)
+        return self.linear(x)
+
+
+Positional = lambda _, emsize: _PositionalEncoding(d_model=emsize)
+
+
+class CanEmb(nn.Embedding):
+    def __init__(self, num_features, num_embeddings: int, embedding_dim: int, *args, **kwargs):
+        assert embedding_dim % num_features == 0
+        embedding_dim = embedding_dim // num_features
+        super().__init__(num_embeddings, embedding_dim, *args, **kwargs)
+
+    def forward(self, x):
+        x = super().forward(x)
+        return x.view(*x.shape[:-2], -1)
+
+def get_Canonical(num_classes):
+    return lambda num_features, emsize: CanEmb(num_features, num_classes, emsize)
+
+def get_Embedding(num_embs_per_feature=100):
+    return lambda num_features, emsize: EmbeddingEncoder(num_features, emsize, num_embs=num_embs_per_feature)

prior-fitting/losses.py

@@ -0,0 +1,12 @@
+import torch
+from torch import nn
+
+
+class ScaledSoftmaxCE(nn.Module):
+    def forward(self, x, label):
+        logits = x[..., :-10]
+        temp_scales = x[..., -10:]
+
+
+
+        logprobs = logits.softmax(-1)

prior-fitting/mcmc_svi_transformer_on_bayesian.py

@@ -0,0 +1,443 @@
+import scipy.stats as st
+from train import Losses
+import argparse
+
+import os
+
+from tqdm import tqdm
+import time
+
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+
+import pyro
+import pyro.distributions as dist
+from pyro.nn import PyroModule, PyroSample
+import torch.nn as nn
+from pyro.infer.autoguide import AutoDiagonalNormal
+from pyro.infer import SVI, Trace_ELBO, Predictive, MCMC, NUTS
+from pyro import infer
+import matplotlib.gridspec as gridspec
+import os.path
+import glob
+from train import train, get_weighted_single_eval_pos_sampler
+import priors
+import encoders
+from pyro.infer import SVGD, RBFSteinKernel
+
+class CausalModel(PyroModule):
+    def __init__(self, model_spec, device='cuda'):
+        super().__init__()
+
+        self.device = device
+        self.num_features = model_spec['num_features']
+
+        mu, sigma = torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)
+
+        self.fc1 = PyroModule[nn.Linear](self.num_features, model_spec['embed'])
+        self.drop = pyro.sample('drop', dist.Categorical(probs=torch.tensor([0.5, 0.5]).expand([model_spec['embed'], self.num_features, 2]))).float()
+        self.fc1.weight = PyroSample(dist.Normal(mu, 0.0000001+self.drop).expand([model_spec['embed'], self.num_features]).to_event(2))
+        self.fc1.bias = PyroSample(dist.Normal(mu, sigma).expand([model_spec['embed']]).to_event(1))
+
+        self.fc2 = PyroModule[nn.Linear](model_spec['embed'], 2)
+        self.fc2.weight = PyroSample(dist.Normal(mu, sigma).expand([2, model_spec['embed']]).to_event(2))
+        self.fc2.bias = PyroSample(dist.Normal(mu, sigma).expand([2]).to_event(1))
+
+        self.model = torch.nn.Sequential(self.fc1, self.fc2)
+
+        self.to(self.device)
+
+    def forward(self, x=None, y=None, seq_len=1):
+        if x is None:
+            with pyro.plate("x_plate", seq_len):
+                d_ = dist.Normal(torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)).expand(
+                    [self.num_features]).to_event(1)
+                x = pyro.sample("x", d_)
+
+        out = self.model(x)
+        mu = out.squeeze()
+        softmax = torch.nn.Softmax(dim=1)
+        # sigma = pyro.sample("sigma", dist.Uniform(torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)))
+        with pyro.plate("data", out.shape[0]):
+            # d_ = dist.Normal(mu, sigma)
+            # obs = pyro.sample("obs", d_, obs=y)
+            s = softmax(mu)
+            obs = pyro.sample('obs', dist.Categorical(probs=s), obs=y).float()
+
+        return x, obs
+
+class BayesianModel(PyroModule):
+    def __init__(self, model_spec, device='cuda'):
+        super().__init__()
+
+        self.device = device
+        self.num_features = model_spec['num_features']
+
+        mu, sigma = torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)
+
+        self.fc1 = PyroModule[nn.Linear](self.num_features, model_spec['embed'])
+        self.fc1.weight = PyroSample(
+            dist.Normal(mu, sigma).expand([model_spec['embed'], self.num_features]).to_event(2))
+        self.fc1.bias = PyroSample(dist.Normal(mu, sigma).expand([model_spec['embed']]).to_event(1))
+
+        self.fc2 = PyroModule[nn.Linear](model_spec['embed'], 2)
+        self.fc2.weight = PyroSample(dist.Normal(mu, sigma).expand([2, model_spec['embed']]).to_event(2))
+        self.fc2.bias = PyroSample(dist.Normal(mu, sigma).expand([2]).to_event(1))
+
+        self.model = torch.nn.Sequential(self.fc1, self.fc2)
+
+        self.to(self.device)
+
+    def forward(self, x=None, y=None, seq_len=1):
+        if x is None:
+            with pyro.plate("x_plate", seq_len):
+                d_ = dist.Normal(torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)).expand(
+                    [self.num_features]).to_event(1)
+                x = pyro.sample("x", d_)
+
+        out = self.model(x)
+        mu = out.squeeze()
+        softmax = torch.nn.Softmax(dim=1)
+        # sigma = pyro.sample("sigma", dist.Uniform(torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)))
+        with pyro.plate("data", out.shape[0]):
+            # d_ = dist.Normal(mu, sigma)
+            # obs = pyro.sample("obs", d_, obs=y)
+            s = softmax(mu)
+            obs = pyro.sample('obs', dist.Categorical(probs=s), obs=y).float()
+
+        return x, obs
+
+
+def get_transformer_config(model_spec):
+    return {'lr': 2.006434218345026e-05
+        , 'epochs': 400
+        , 'dropout': 0.0
+        , 'emsize': 256
+        , 'batch_size': 256
+        , 'nlayers': 5
+        , 'num_outputs': 1
+        , 'num_features': model_spec['num_features']
+        , 'steps_per_epoch': 100
+        , 'nhead': 4
+        , 'dropout': 0.0
+        , 'seq_len': model_spec['seq_len']
+        , 'nhid_factor': 2}
+
+
+def get_model(model_generator, config, should_train=True, device='cuda'):
+    epochs = 0 if not should_train else config['epochs']
+
+    model = train(priors.pyro.DataLoader
+                  , Losses.bce
+                  , encoders.Linear
+                  , emsize=config['emsize']
+                  , nhead=config['nhead']
+                  , y_encoder_generator=encoders.Linear
+                  , pos_encoder_generator=None
+                  , batch_size=config['batch_size']
+                  , nlayers=config['nlayers']
+                  , nhid=config['emsize'] * config['nhid_factor']
+                  , epochs=epochs
+                  , warmup_epochs=config['epochs'] // 4
+                  , bptt=config['seq_len']
+                  , gpu_device=device
+                  , dropout=config['dropout']
+                  , steps_per_epoch=config['steps_per_epoch']
+                  , single_eval_pos_gen=get_weighted_single_eval_pos_sampler(100)
+                  , extra_prior_kwargs_dict={
+            'num_outputs': config['num_outputs']
+            , 'num_features': config['num_features']
+            , 'canonical_args': None
+            , 'fuse_x_y': False
+            , 'model': model_generator
+        }
+                  , lr=config['lr']
+                  , verbose=True)
+
+    return model
+
+
+
+def plot_features(data, targets):
+    fig2 = plt.figure(constrained_layout=True, figsize=(12, 12))
+    spec2 = gridspec.GridSpec(ncols=data.shape[1], nrows=data.shape[1], figure=fig2)
+    for d in range(0, data.shape[1]):
+        for d2 in range(0, data.shape[1]):
+            sub_ax = fig2.add_subplot(spec2[d, d2])
+            sub_ax.scatter(data[:, d].detach().cpu().numpy(), data[:, d2].detach().cpu().numpy(),
+                           c=targets[:].detach().cpu().numpy())
+
+
+def evaluate_preds(preds, y_test):
+    preds_hard = preds['obs'] > 0.5  # TODO: 0.5 or 0
+    acc = (preds_hard == y_test).float().mean()
+    means = preds_hard.float().mean(axis=0)
+
+    # var = preds['obs'].var(axis=0)
+    nll = nn.BCELoss()(means.float(), y_test.float())
+    mse = Losses.mse(means, y_test).mean()
+
+    return acc, nll, mse
+
+
+def load_results(path, task='steps'):
+    results_nll = []
+    results_acc = []
+    times = []
+    samples_list = []
+
+    files = glob.glob(f'/home/hollmann/prior-fitting/{path}_*.npy')
+    for file in files:
+        print(file)
+        with open(file, 'rb') as f:
+            if task == 'steps':
+                nll, acc, elapsed = np.load(f, allow_pickle=True)
+                samples_list += [file]
+            else:
+                samples, nll, acc, elapsed = np.load(f, allow_pickle=True)
+                samples_list += [samples]
+            times += [elapsed]
+            results_nll += [nll]
+            results_acc += [acc]
+    results_acc = np.array(results_acc)
+    results_nll = np.array(results_nll)
+    times = np.array(times)
+    files = np.array(files)
+    samples = np.array(samples_list)
+    means = np.array([compute_mean_and_conf_interval(results_nll[n, :])[0] for n in range(0, results_nll.shape[0])])
+    conf = np.array([compute_mean_and_conf_interval(results_nll[n, :])[1] for n in range(0, results_nll.shape[0])])
+
+    if task == 'steps':
+        sorter = np.argsort(times, axis=0)
+    else:
+        sorter = np.argsort(samples, axis=0)
+
+    results_nll, results_acc, times, files, samples, means, conf = results_nll[sorter], results_acc[sorter], times[sorter], files[sorter], samples[sorter], means[sorter], conf[sorter]
+
+    return files, times, samples, means, conf
+
+def plot_with_confidence_intervals(ax_or_pyplot, x, mean, confidence, **common_kwargs):
+    ax_or_pyplot.plot(x,mean,**common_kwargs)
+    if 'label' in common_kwargs:
+        common_kwargs.pop('label')
+    if 'marker' in common_kwargs:
+        common_kwargs.pop('marker')
+    ax_or_pyplot.fill_between(x, (mean-confidence), (mean+confidence), alpha=.1, **common_kwargs)
+
+
+def compute_mean_and_conf_interval(accuracies, confidence=.95):
+    accuracies = np.array(accuracies)
+    n = len(accuracies)
+    m, se = np.mean(accuracies), st.sem(accuracies)
+    h = se * st.t.ppf((1 + confidence) / 2., n - 1)
+    return m, h
+
+
+def generate_toy_data(model, bptt, device='cpu'):
+    n_samples = 100
+    X_list, y_list = [], []
+    torch.manual_seed(0)
+    for _ in range(0, n_samples):
+        X_sample, y_sample = model(seq_len=bptt)
+        X_list += [X_sample]
+        y_list += [y_sample]
+    X = torch.stack(X_list, 0)
+    y = torch.stack(y_list, 0)
+    # y = (y > 0).float()
+
+    return X.to(device), y.to(device)
+
+
+
+def eval_svi(X, y, device, model_sampler, training_samples_n, num_train_steps, num_pred_samples, lr=1e-3, num_particles=1, svgd=False):
+    X_test, y_test = X[:, training_samples_n:], y[:, training_samples_n:]
+    X_train, y_train = X[:, 0:training_samples_n], y[:, 0:training_samples_n]
+
+    nll_list = []
+    acc_list = []
+    for sample_id in tqdm(list(range(0, X_test.shape[0]))):
+        model = model_sampler()
+        guide = AutoDiagonalNormal(model).to(device)
+        adam = pyro.optim.Adam({"lr": lr})
+        svi = SVI(model, guide, adam, loss=Trace_ELBO(num_particles=num_particles))
+
+        if svgd:
+            kernel = RBFSteinKernel()
+            svi = SVGD(model, kernel, adam, num_particles=50, max_plate_nesting=0)
+
+        pyro.clear_param_store()
+
+        X_test_sample, y_test_sample, X_train_sample, y_train_sample = X_test[sample_id], y_test[sample_id], X_train[
+            sample_id], y_train[sample_id]
+
+        acc, nll, mse = 0.0, 0.0, 0.0
+        # bar = tqdm(list(range(num_train_steps)))
+        bar = list(range(num_train_steps))
+        for epoch in bar:
+            loss = svi.step(X_train_sample, y_train_sample)
+            # if epoch % 100 == 1:
+            #    bar.set_postfix(loss=f'{loss / X_train_sample.shape[0]:.3f}', test_nll=f'{nll:.3f}', test_acc=f'{acc:.3f}')
+
+        predictive = Predictive(model, guide=guide, num_samples=num_pred_samples)
+        preds = predictive(X_test_sample)
+        acc, nll, mse = evaluate_preds(preds, y_test_sample)
+        nll_list += [nll.detach().cpu().numpy()]
+        acc_list += [acc.detach().cpu().numpy()]
+
+    return np.array(nll_list), np.array(acc_list)
+
+
+def eval_mcmc(X, y, device, model_sampler, training_samples_n, warmup_steps, num_pred_samples):
+    X_test, y_test = X[:, training_samples_n:].to(device), y[:, training_samples_n:].to(device)
+    X_train, y_train = X[:, 0:training_samples_n].to(device), y[:, 0:training_samples_n].to(device)
+
+    acc_list, nll_list = [], []
+    for sample_id in tqdm(list(range(0, X_test.shape[0]))):
+        X_test_sample, y_test_sample, X_train_sample, y_train_sample = X_test[sample_id], y_test[sample_id], X_train[
+            sample_id], y_train[sample_id]
+
+        model = model_sampler()
+        mcmc = MCMC(NUTS(model), num_samples=num_pred_samples, num_chains=1, disable_progbar=True,
+                    warmup_steps=warmup_steps, mp_context="fork")
+        mcmc.run(X_train_sample, y_train_sample)
+        preds = infer.mcmc.util.predictive(model, mcmc.get_samples(), X_test_sample, None)
+        acc, nll, mse = evaluate_preds(preds, y_test_sample)
+        nll_list += [nll.detach().cpu().numpy()]
+        acc_list += [acc.detach().cpu().numpy()]
+
+    return np.array(nll_list), np.array(acc_list)
+
+
+def eval_transformer(X, y, device, model, training_samples_n):
+    X_sample, y_sample = X.transpose(0, 1), y.transpose(0, 1).float()
+    bs = 1
+    samples = []
+    for i in range(0, X_sample.shape[1] // bs):
+        samples += [(X_sample[:, bs * i:bs * (i + 1)], y_sample[:, bs * i:bs * (i + 1)])]
+
+    mean = X_sample[:training_samples_n].mean(0)
+    std = X_sample[:training_samples_n].std(0) + .000001
+    X_sample = (X_sample - mean) / std
+
+    start = time.time()
+    output = torch.cat(
+        [model.to(device)((X_sample_chunk, y_sample_chunk), single_eval_pos=training_samples_n).squeeze(-1) for
+         (X_sample_chunk, y_sample_chunk) in samples], 1)
+    elapsed = time.time() - start
+
+    output = output.detach().cpu()
+    acc = ((torch.sigmoid(output) > 0.5) == y_sample[training_samples_n:].cpu().bool()).float().mean(axis=0)
+    nll = nn.BCELoss(reduction='none')(torch.sigmoid(output.float()), y_sample[training_samples_n:].cpu().float()).mean(
+        axis=0)
+    return acc, nll, elapsed
+
+
+def training_steps(method, X, y, model_spec, device='cpu', path_interfix='', overwrite=False):
+    training_samples_n = 100
+    for s in [2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]:
+        path = f'/home/hollmann/prior-fitting/{path_interfix}/results_{method}_training_steps_{s}.npy'
+        if (os.path.isfile(path)) and not overwrite:
+            print(f'already done {s}')
+            continue
+
+        start = time.time()
+        if method == 'svi':
+            nll, acc = eval_svi(X, y, device, model_spec, training_samples_n, num_train_steps=s, num_pred_samples=s, svgd=False)
+        elif method == 'svgd':
+            nll, acc = eval_svi(X, y, device, model_spec, training_samples_n, num_train_steps=s, num_pred_samples=s, svgd=True)
+        elif method == 'mcmc':
+            nll, acc = eval_mcmc(X, y, device, model_spec, training_samples_n, warmup_steps=s, num_pred_samples=s)
+        elapsed = time.time() - start
+
+        print(s)
+        print('NLL ', compute_mean_and_conf_interval(nll))
+        print('ACC ', compute_mean_and_conf_interval(acc))
+        print('TIME ', elapsed)
+
+        with open(path, 'wb') as f:
+            np.save(f, (np.array(nll), np.array(acc), elapsed))
+
+        print(f'Saved results at {path}')
+
+
+def training_samples(method, X, y, model_spec, evaluation_points, steps = None, device='cpu', path_interfix='', overwrite=False):
+    num_pred_samples_mcmc = steps if steps else 512
+    warmup_steps = steps if steps else 512
+
+    num_pred_samples_svi = steps if steps else 1024
+    num_train_steps = steps if steps else 1024
+
+    num_pred_samples = num_pred_samples_svi if method == 'svi' else num_pred_samples_mcmc
+
+    for training_samples_n in evaluation_points:
+        path = f'/home/hollmann/prior-fitting/{path_interfix}/results_{method}_{num_pred_samples}_training_samples_{training_samples_n}.npy'
+        if (os.path.isfile(path)) and not overwrite:
+            print(f'already done {training_samples_n}')
+            continue
+
+        start = time.time()
+        if method == 'svi':
+            nll, acc = eval_svi(X, y, device, model_spec, training_samples_n, num_train_steps=num_train_steps, num_pred_samples=num_pred_samples)
+        elif method == 'svgd':
+            nll, acc = eval_svi(X, y, device, model_spec, training_samples_n, num_train_steps=num_train_steps, num_pred_samples=num_pred_samples, svgd=True)
+        elif method == 'mcmc':
+            nll, acc = eval_mcmc(X, y, device, model_spec, training_samples_n, warmup_steps=warmup_steps, num_pred_samples=num_pred_samples)
+        elapsed = time.time() - start
+
+        print('NLL ', compute_mean_and_conf_interval(nll))
+        print('ACC ', compute_mean_and_conf_interval(acc))
+        print('TIME ', elapsed)
+
+        with open(path, 'wb') as f:
+            np.save(f, (training_samples_n, np.array(nll), np.array(acc), elapsed))
+
+### MAIN
+def get_default_model_spec(size):
+    bptt = 300
+
+    if size == 'big':
+        num_features = 8
+        embed = 64
+        nlayers = 2
+    elif size == 'small':
+        num_features = 3
+        embed = 5
+        nlayers = 2
+    else:
+        num_features = int(size.split("_")[0])
+        embed = int(size.split("_")[1])
+        nlayers = int(size.split("_")[2])
+
+    return {'nlayers': nlayers, 'embed': embed, 'num_features': num_features, "seq_len": bptt}
+
+def get_default_evaluation_points():
+    return list(range(2, 100, 5))
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--solver', default='svi', type=str)
+    parser.add_argument('--task', default='steps', type=str)
+    parser.add_argument('--model_size', default='small', type=str)
+
+    args = parser.parse_args()
+
+    model_spec = get_default_model_spec(args.model_size)
+    evaluation_points = get_default_evaluation_points()
+    device = 'cuda:0' if args.solver == 'svi' else 'cpu'
+
+    torch.manual_seed(0)
+    test_model = BayesianModel(model_spec, device=device)
+
+    X, y = generate_toy_data(test_model, model_spec['seq_len'])
+    model_sampler = lambda: BayesianModel(model_spec, device=device)
+
+    if args.task == 'steps':
+        training_steps(args.solver, X, y, model_sampler, device=device,
+                       path_interfix=f'results/timing_{args.model_size}_model', svgd=args.svgd)
+    elif args.task == 'samples':
+        training_samples(args.solver, X, y, model_sampler, evaluation_points, device=device,
+                       path_interfix=f'results/timing_{args.model_size}_model', svgd=args.svgd)
+
+
+

prior-fitting/notebooks/BayesianModels_And_Custom_Pyro_Modules.ipynb

@@ -0,0 +1,524 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 56,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import numpy as np\n",
+    "\n",
+    "import priors\n",
+    "from train import train, get_weighted_single_eval_pos_sampler\n",
+    "import encoders\n",
+    "import positional_encodings\n",
+    "import utils\n",
+    "import bar_distribution\n",
+    "import decoders\n",
+    "from datasets import *\n",
+    "import os\n",
+    "\n",
+    "from tqdm import tqdm\n",
+    "import time\n",
+    "\n",
+    "import torch\n",
+    "import pandas as pd\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import torch.nn as nn\n",
+    "import os.path\n",
+    "import glob\n",
+    "\n",
+    "from mcmc_svi_transformer_on_bayesian import get_model, get_default_model_spec, generate_toy_data, load_results, plot_with_confidence_intervals, training_steps, training_samples, get_default_evaluation_points, compute_mean_and_conf_interval, eval_transformer\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# %load_ext autoreload\n",
+    "\n",
+    "# %autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## DEFINE A PRIOR MODEL ##\n",
+    "# We define a Bayesian Model as a prior for all methods\n",
+    "# This can be replaced by other models that inherit from PyroModule.\n",
+    "class BayesianModel(PyroModule):\n",
+    "    def __init__(self, model_spec, device='cuda'):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        self.device = device\n",
+    "        self.num_features = model_spec['num_features']\n",
+    "\n",
+    "        mu, sigma = torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)\n",
+    "\n",
+    "        self.fc1 = PyroModule[nn.Linear](self.num_features, model_spec['embed'])\n",
+    "        self.fc1.weight = PyroSample(\n",
+    "            dist.Normal(mu, sigma).expand([model_spec['embed'], self.num_features]).to_event(2))\n",
+    "        self.fc1.bias = PyroSample(dist.Normal(mu, sigma).expand([model_spec['embed']]).to_event(1))\n",
+    "\n",
+    "        self.fc2 = PyroModule[nn.Linear](model_spec['embed'], 2)\n",
+    "        self.fc2.weight = PyroSample(dist.Normal(mu, sigma).expand([2, model_spec['embed']]).to_event(2))\n",
+    "        self.fc2.bias = PyroSample(dist.Normal(mu, sigma).expand([2]).to_event(1))\n",
+    "\n",
+    "        self.model = torch.nn.Sequential(self.fc1, self.fc2)\n",
+    "\n",
+    "        self.to(self.device)\n",
+    "\n",
+    "    def forward(self, x=None, y=None, seq_len=1):\n",
+    "        if x is None:\n",
+    "            with pyro.plate(\"x_plate\", seq_len):\n",
+    "                d_ = dist.Normal(torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)).expand(\n",
+    "                    [self.num_features]).to_event(1)\n",
+    "                x = pyro.sample(\"x\", d_)\n",
+    "\n",
+    "        out = self.model(x)\n",
+    "        mu = out.squeeze()\n",
+    "        softmax = torch.nn.Softmax(dim=1)\n",
+    "        with pyro.plate(\"data\", out.shape[0]):\n",
+    "            s = softmax(mu)\n",
+    "            obs = pyro.sample('obs', dist.Categorical(probs=s), obs=y).float()\n",
+    "\n",
+    "        return x, obs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 69,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "results_directory = 'results' # Where to save results\n",
+    "model_spec_size = 'small' # Size of the BNN model to evaluate, also try big\n",
+    "bptt = 100 # Number of samples in each dataset\n",
+    "\n",
+    "# Training samples seen after which to evaluate the methods\n",
+    "evaluation_points = [2, 7, 12, 17, 22, 27, 32, 37, 42, 47, 52, 57, 62, 67, 72, 77, 82, 87, 92]\n",
+    "\n",
+    "# Function which generates a model from the prior\n",
+    "model_sampler = lambda : BayesianModel(get_default_model_spec(model_spec_size), device = device)\n",
+    "\n",
+    "global_results = {} # Dict in which to save results\n",
+    "task = 'samples' # Task to evaluate, only option is samples, keep fixed"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!mkdir {results_directory}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "heading_collapsed": true
+   },
+   "source": [
+    "### Evaluate SVI and MCMC"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {
+    "hidden": true
+   },
+   "outputs": [],
+   "source": [
+    "method = 'svi'\n",
+    "steps = 1\n",
+    "device = 'cuda'\n",
+    "path_interfix = f'{results_directory}/timing_{model_spec_size}_model_test'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "metadata": {
+    "hidden": true
+   },
+   "outputs": [],
+   "source": [
+    "!mkdir {path_interfix}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "metadata": {
+    "hidden": true
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 100/100 [00:02<00:00, 37.13it/s]\n",
+      "/home/hollmann/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/numpy/lib/npyio.py:528: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.\n",
+      "  arr = np.asanyarray(arr)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "NLL  (51.540817, 1.832436208065078)\n",
+      "ACC  (0.48459178, 0.01832436154844232)\n",
+      "TIME  2.6950523853302\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 100/100 [00:02<00:00, 35.89it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "NLL  (48.569893, 1.8696300575034437)\n",
+      "ACC  (0.51430106, 0.01869630134377999)\n",
+      "TIME  2.788970708847046\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 100/100 [00:03<00:00, 31.80it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "NLL  (51.034092, 1.807273770560027)\n",
+      "ACC  (0.48965907, 0.018072737823868815)\n",
+      "TIME  3.1472866535186768\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 100/100 [00:02<00:00, 38.48it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "NLL  (50.216866, 2.0121896389094833)\n",
+      "ACC  (0.4978313, 0.02012189562034928)\n",
+      "TIME  2.600956439971924\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      " 55%|█████▌    | 55/100 [00:01<00:01, 38.41it/s]\n"
+     ]
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
+      "\u001b[0;32m/tmp/ipykernel_9449/1948451174.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgenerate_toy_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtest_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbptt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m training_samples(method\n\u001b[0m\u001b[1;32m      5\u001b[0m                  \u001b[0;34m,\u001b[0m \u001b[0mX\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m                  \u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/prior-fitting/mcmc_svi_transformer_on_bayesian.py\u001b[0m in \u001b[0;36mtraining_samples\u001b[0;34m(method, X, y, model_spec, evaluation_points, steps, device, path_interfix, overwrite)\u001b[0m\n\u001b[1;32m    379\u001b[0m         \u001b[0mstart\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    380\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mmethod\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'svi'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 381\u001b[0;31m             \u001b[0mnll\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0macc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0meval_svi\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel_spec\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtraining_samples_n\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_train_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnum_train_steps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_pred_samples\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnum_pred_samples\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    382\u001b[0m         \u001b[0;32melif\u001b[0m \u001b[0mmethod\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'svgd'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    383\u001b[0m             \u001b[0mnll\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0macc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0meval_svi\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel_spec\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtraining_samples_n\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_train_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnum_train_steps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_pred_samples\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnum_pred_samples\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msvgd\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/prior-fitting/mcmc_svi_transformer_on_bayesian.py\u001b[0m in \u001b[0;36meval_svi\u001b[0;34m(X, y, device, model_sampler, training_samples_n, num_train_steps, num_pred_samples, lr, num_particles, svgd)\u001b[0m\n\u001b[1;32m    281\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    282\u001b[0m         \u001b[0mpredictive\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mPredictive\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mguide\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mguide\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_samples\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnum_pred_samples\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 283\u001b[0;31m         \u001b[0mpreds\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpredictive\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX_test_sample\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    284\u001b[0m         \u001b[0macc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnll\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmse\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mevaluate_preds\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpreds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_test_sample\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    285\u001b[0m         \u001b[0mnll_list\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mnll\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1049\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1050\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1051\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1052\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1053\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/infer/predictive.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    271\u001b[0m                 \u001b[0mmodel_kwargs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    272\u001b[0m             )\n\u001b[0;32m--> 273\u001b[0;31m         return _predictive(\n\u001b[0m\u001b[1;32m    274\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    275\u001b[0m             \u001b[0mposterior_samples\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/infer/predictive.py\u001b[0m in \u001b[0;36m_predictive\u001b[0;34m(model, posterior_samples, num_samples, return_sites, return_trace, parallel, model_args, model_kwargs)\u001b[0m\n\u001b[1;32m    125\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    126\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mparallel\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 127\u001b[0;31m         return _predictive_sequential(\n\u001b[0m\u001b[1;32m    128\u001b[0m             \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    129\u001b[0m             \u001b[0mposterior_samples\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/infer/predictive.py\u001b[0m in \u001b[0;36m_predictive_sequential\u001b[0;34m(model, posterior_samples, model_args, model_kwargs, num_samples, return_site_shapes, return_trace)\u001b[0m\n\u001b[1;32m     46\u001b[0m     ]\n\u001b[1;32m     47\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnum_samples\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 48\u001b[0;31m         trace = poutine.trace(poutine.condition(model, samples[i])).get_trace(\n\u001b[0m\u001b[1;32m     49\u001b[0m             \u001b[0;34m*\u001b[0m\u001b[0mmodel_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mmodel_kwargs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     50\u001b[0m         )\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/poutine/trace_messenger.py\u001b[0m in \u001b[0;36mget_trace\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    196\u001b[0m         \u001b[0mCalls\u001b[0m \u001b[0mthis\u001b[0m \u001b[0mpoutine\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mreturns\u001b[0m \u001b[0mits\u001b[0m \u001b[0mtrace\u001b[0m \u001b[0minstead\u001b[0m \u001b[0mof\u001b[0m \u001b[0mthe\u001b[0m \u001b[0mfunction\u001b[0m\u001b[0;31m'\u001b[0m\u001b[0ms\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    197\u001b[0m         \"\"\"\n\u001b[0;32m--> 198\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    199\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmsngr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_trace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/poutine/trace_messenger.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    172\u001b[0m             )\n\u001b[1;32m    173\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 174\u001b[0;31m                 \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    175\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mValueError\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    176\u001b[0m                 \u001b[0mexc_type\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexc_value\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtraceback\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexc_info\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/poutine/messenger.py\u001b[0m in \u001b[0;36m_context_wrap\u001b[0;34m(context, fn, *args, **kwargs)\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_context_wrap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcontext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mcontext\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 12\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     13\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/torch/autograd/grad_mode.py\u001b[0m in \u001b[0;36mdecorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     26\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__class__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 28\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     29\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mF\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/poutine/messenger.py\u001b[0m in \u001b[0;36m_context_wrap\u001b[0;34m(context, fn, *args, **kwargs)\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_context_wrap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcontext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mcontext\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 12\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     13\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/nn/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    424\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    425\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_pyro_context\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 426\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    427\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    428\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__getattr__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1049\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1050\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1051\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1052\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1053\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m/tmp/ipykernel_9449/3309204952.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x, y, seq_len)\u001b[0m\n\u001b[1;32m     28\u001b[0m                 \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpyro\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msample\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"x\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md_\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 30\u001b[0;31m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     31\u001b[0m         \u001b[0mmu\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     32\u001b[0m         \u001b[0msoftmax\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mSoftmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1049\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1050\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1051\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1052\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1053\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    137\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    138\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 139\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    140\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    141\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/nn/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    424\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    425\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_pyro_context\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 426\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    427\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    428\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__getattr__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1049\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1050\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1051\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1052\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1053\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/torch/nn/modules/linear.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m     94\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     95\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 96\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlinear\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     97\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     98\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mextra_repr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/nn/module.py\u001b[0m in \u001b[0;36m__getattr__\u001b[0;34m(self, name)\u001b[0m\n\u001b[1;32m    477\u001b[0m                         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprior\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"sample\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# if not a distribution\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    478\u001b[0m                             \u001b[0mprior\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprior\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 479\u001b[0;31m                         \u001b[0mvalue\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpyro\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msample\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfullname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprior\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    480\u001b[0m                         \u001b[0mcontext\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfullname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    481\u001b[0m                     \u001b[0;32mreturn\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/primitives.py\u001b[0m in \u001b[0;36msample\u001b[0;34m(name, fn, *args, **kwargs)\u001b[0m\n\u001b[1;32m    162\u001b[0m         }\n\u001b[1;32m    163\u001b[0m         \u001b[0;31m# apply the stack and return its return value\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 164\u001b[0;31m         \u001b[0mapply_stack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    165\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mmsg\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"value\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    166\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/poutine/runtime.py\u001b[0m in \u001b[0;36mapply_stack\u001b[0;34m(initial_msg)\u001b[0m\n\u001b[1;32m    210\u001b[0m         \u001b[0mpointer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpointer\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    211\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 212\u001b[0;31m         \u001b[0mframe\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_process_message\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    213\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    214\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mmsg\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"stop\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/pyro/poutine/messenger.py\u001b[0m in \u001b[0;36m_process_message\u001b[0;34m(self, msg)\u001b[0m\n\u001b[1;32m    139\u001b[0m         \u001b[0mon\u001b[0m \u001b[0mmessage\u001b[0m \u001b[0mtype\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mThe\u001b[0m \u001b[0mmessage\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0mupdated\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mplace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    140\u001b[0m         \"\"\"\n\u001b[0;32m--> 141\u001b[0;31m         \u001b[0mmethod\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"_pyro_{}\"\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"type\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    142\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mmethod\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    143\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mmethod\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     ]
+    }
+   ],
+   "source": [
+    "X, y = generate_toy_data(test_model, bptt, device)\n",
+    "\n",
+    "training_samples(method\n",
+    "                 , X\n",
+    "                 , y\n",
+    "                 , model_sampler\n",
+    "                 , evaluation_points\n",
+    "                 , steps=steps\n",
+    "                 , device=device\n",
+    "                 , path_interfix=path_interfix)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "heading_collapsed": true
+   },
+   "source": [
+    "### Training Transformer on Prior (Skip this step to reuse results)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 41,
+   "metadata": {
+    "hidden": true
+   },
+   "outputs": [],
+   "source": [
+    "device = 'cuda'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 49,
+   "metadata": {
+    "hidden": true
+   },
+   "outputs": [],
+   "source": [
+    "config = {'lr': 2.006434218345026e-05\n",
+    " , 'epochs': 160\n",
+    " , 'dropout': 0.0\n",
+    " , 'emsize': 256\n",
+    " , 'batch_size': 256\n",
+    " , 'nlayers': 5\n",
+    " , 'num_outputs': 1\n",
+    " , 'num_features': model_spec['num_features']\n",
+    " , 'steps_per_epoch': 100\n",
+    " , 'nhead': 4\n",
+    " , 'dropout': 0.0\n",
+    " , 'seq_len': model_spec['seq_len']\n",
+    " , 'nhid_factor': 2}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 51,
+   "metadata": {
+    "hidden": true
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using cuda device\n",
+      "DataLoader.__dict__ {'num_steps': 100, 'fuse_x_y': False, 'get_batch_kwargs': {'batch_size': 256, 'seq_len': 300, 'num_outputs': 1, 'num_features': 3, 'canonical_args': None, 'model': <function <lambda> at 0x7f6f42f49f70>}, 'num_features': 3, 'num_outputs': 1}\n"
+     ]
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
+      "\u001b[0;32m/tmp/ipykernel_9449/1283571267.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtransformer_model\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel_sampler\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshould_train\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mmodel_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresults_directory\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34mf'bayesian_models_transformer_checkpoint_{model_spec_size}_epochs_'\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'epochs'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;34m'.cpkt'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtransformer_model\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate_dict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/prior-fitting/mcmc_svi_transformer_on_bayesian.py\u001b[0m in \u001b[0;36mget_model\u001b[0;34m(model_generator, config, should_train, device)\u001b[0m\n\u001b[1;32m    129\u001b[0m     \u001b[0mepochs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mshould_train\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'epochs'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    130\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 131\u001b[0;31m     model = train(priors.pyro.DataLoader\n\u001b[0m\u001b[1;32m    132\u001b[0m                   \u001b[0;34m,\u001b[0m \u001b[0mLosses\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbce\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    133\u001b[0m                   \u001b[0;34m,\u001b[0m \u001b[0mencoders\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mLinear\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/prior-fitting/train.py\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(priordataloader_class, criterion, encoder_generator, emsize, nhid, nlayers, nhead, dropout, epochs, steps_per_epoch, batch_size, bptt, lr, warmup_epochs, input_normalization, y_encoder_generator, pos_encoder_generator, decoder, extra_prior_kwargs_dict, scheduler, load_weights_from_this_state_dict, validation_period, single_eval_pos_gen, gpu_device, aggregate_k_gradients, verbose)\u001b[0m\n\u001b[1;32m    116\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepochs\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    117\u001b[0m         \u001b[0mepoch_start_time\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 118\u001b[0;31m         \u001b[0mtotal_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtotal_positional_losses\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtime_to_get_batch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mforward_time\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstep_time\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    119\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdl\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'validate'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0mvalidation_period\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    120\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mno_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/prior-fitting/train.py\u001b[0m in \u001b[0;36mtrain\u001b[0;34m()\u001b[0m\n\u001b[1;32m     93\u001b[0m             \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     94\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mbatch\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0maggregate_k_gradients\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0maggregate_k_gradients\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 95\u001b[0;31m                 \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mutils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclip_grad_norm_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparameters\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1.\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     96\u001b[0m                 \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     97\u001b[0m                 \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/miniconda3/envs/prior-fitting/lib/python3.9/site-packages/torch/nn/utils/clip_grad.py\u001b[0m in \u001b[0;36mclip_grad_norm_\u001b[0;34m(parameters, max_norm, norm_type, error_if_nonfinite)\u001b[0m\n\u001b[1;32m     41\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     42\u001b[0m         \u001b[0mtotal_norm\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrad\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnorm_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mp\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mparameters\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnorm_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 43\u001b[0;31m     \u001b[0;32mif\u001b[0m \u001b[0mtotal_norm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misnan\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mtotal_norm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misinf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     44\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0merror_if_nonfinite\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     45\u001b[0m             raise RuntimeError(\n",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     ]
+    }
+   ],
+   "source": [
+    "transformer_model = get_model(model_sampler, config, should_train = True)\n",
+    "model_path = os.path.join(results_directory, f'bayesian_models_transformer_checkpoint_{model_spec_size}_epochs_'+config['epochs']+'.cpkt')\n",
+    "torch.save((transformer_model[2].state_dict(), None), model_path)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Evaluating Transformer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 52,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using cuda device\n",
+      "DataLoader.__dict__ {'num_steps': 100, 'fuse_x_y': False, 'get_batch_kwargs': {'batch_size': 256, 'seq_len': 300, 'num_outputs': 1, 'num_features': 3, 'canonical_args': None, 'model': <function <lambda> at 0x7f6f42f49f70>}, 'num_features': 3, 'num_outputs': 1}\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "<All keys matched successfully>"
+      ]
+     },
+     "execution_count": 52,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "loaded_epoch = config['epochs']\n",
+    "transformer_model = get_model(model_sampler, config, should_train = False)\n",
+    "path = os.path.join(results_directory, F'bayesian_models_transformer_checkpoint_{model_spec_size}_epochs_{loaded_epoch}.cpkt')\n",
+    "model_state, optimizer_state = torch.load(path)\n",
+    "transformer_model[2].load_state_dict(model_state)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 57,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "X, y = generate_toy_data(test_model, bptt, device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 73,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "results_acc = []\n",
+    "results_nll = []\n",
+    "transformer_model[2].eval()\n",
+    "for training_samples_n in evaluation_points:\n",
+    "    acc, nll, elapsed = eval_transformer(X, y, model=transformer_model[2], training_samples_n=training_samples_n, device=device)\n",
+    "    results_acc.append(acc)\n",
+    "    results_nll.append(nll)\n",
+    "mean = np.array([compute_mean_and_conf_interval(nll)[0] for nll in results_nll])\n",
+    "conf = np.array([compute_mean_and_conf_interval(nll)[1] for nll in results_nll])\n",
+    "\n",
+    "global_results['transformer'] = (None, np.array(evaluation_points), mean, conf)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Plotting results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 71,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+     ]
+    }
+   ],
+   "source": [
+    "files, times, samples, mean, conf = load_results(f'{results_directory}/timing_{model_size}_model/results_svi_training_{task}', task=task)\n",
+    "global_results['svi'] = (times/100, samples, mean, conf)\n",
+    "files, times, samples, mean, conf = load_results(f'{results_directory}/timing_{model_size}_model/results_mcmc_training_{task}', task=task)\n",
+    "global_results['mcmc'] = (times/100, samples,mean, conf)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 74,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<matplotlib.legend.Legend at 0x7f6f1c2ba7c0>"
+      ]
+     },
+     "execution_count": 74,
+     "metadata": {},
+     "output_type": "execute_result"
+    },
+    {
+     "data": {
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAEoCAYAAAAub0k8AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAABkaUlEQVR4nO3dd5xcV33//9e5ZdrubFXfVe9dtiwXDLaMwYWOIfQaCCGE/L4hgQAB0wkQQgKh+QsJhFACGPgGgw0GG1dwkWTLlizJ6mVX0vY69Zbz++POzM5WrVaz2vZ5Phhm5s6dmbvj1d73nPM55yitNUIIIYSYWYyJPgAhhBBCXHwSAIQQQogZSAKAEEIIMQNJABBCCCFmIAkAQgghxAwkAUAIIYSYgayJeuNZs2bpJUuWTNTbCyGEEFParl27WrXWs8f6/AkLAEuWLGHnzp0T9fZCCCHElKaUOnEhz5cuACGEEGIGkgAghBBCzEASAIQQQogZaMJqAIQQQkwdjuPQ0NBAOp2e6EOZcSKRCPX19di2XdLXlQAghBDinBoaGojH4yxZsgSl1EQfzoyhtaatrY2GhgaWLl1a0teWLgAhhBDnlE6nqa2tlZP/RaaUora2dlxaXiQACCGEGBU5+U+M8frcJQAIIYSYdl70ohfR2dk50YcxqUkNgBBCiGnnrrvumuhDmPSkBUAIIcSUkEgkePGLX8zmzZvZsGED3/ve93jNa15TePz+++/npS99KRDMNtva2jpRhzolTFgLgJ6oNxZCCHFBPvmrZ9h3urukr7luQQUff+n6Eff57W9/y4IFC7jzzjsB6Orq4tZbbyWRSFBWVsZPfvITXvva15b0uKazCWsB0FoigBBCiNHbuHEj99xzDx/84Ad56KGHqKys5KabbuJXv/oVruty55138vKXv3yiD3PKmLAWAF8CgBBCTEnn+qY+XlatWsWuXbu46667+PCHP8wNN9zAa1/7Wr7+9a9TU1PDtm3biMfjE3JsU9GEtgC4nj9Rby+EEGKKOX36NLFYjDe96U28//3v54knnmD79u088cQTfPvb35bm//M0gTUAmqznY5lShyiEEOLc9uzZwwc+8AEMw8C2bb75zW9imiYveclL+K//+i++973vTfQhTilqovriN27ZrB/84+NUl4Un5P2FEEKM3v79+1m7du1EH8aMNdTnr5TapbW+bKyvOaFfv5NOdiLfXgghhJixJjQAOL5D1pU6ACGEEOJim/gAIIWAQgghxEU3oQHAlRYAIYQQYkJMaADwtSbjZGVSICGEEOIim/AxeFnflW4AIYQQ4iKb8AAg3QBCCCHExTdhAcD1gmZ/R0sAEEIIIS62CQsAx9vS+Frj+R5Zz8P3pQ5ACCHE8I4fP86aNWt45zvfyYYNG3jjG9/IPffcw9VXX83KlSt5/PHH6e3t5e1vfzsbN25k06ZN/PznPwegvLycD37wg2zdupUXvOAFPP7442zfvp1ly5Zxxx13AOB5Hu9///sLz/3qV786kT/uuJvQxYCONqdYMTcWdAN4YSKGOVGHI4QQYrR+8yE4u6e0rzlvI9z8+XPudvjwYW6//Xa+9a1vsW3bNn70ox/x8MMPc8cdd/BP//RPrF69msrKSvbsCY6vo6MDgEQiwfbt2/nCF77AK1/5Sj760Y/y+9//nn379vHWt76Vl73sZXzrW9/i2LFjPPnkk1iWRXt7e2l/xklmwgKAMhM8dao3FwBcMq5PxJYAIIQQYnhLly5l48aNAKxfv57rr78epRQbN27k+PHjnDp1ih//+MeF/aurqwEIhULcdNNNQLCscDgcxrbtwvMA7rnnHt797ndjWcGpsaam5iL+ZBffxAUAu5MnTnbwqsvmyIyAQggxlYzim/p4CYf71o8xDKNw3zAMXNfFMAyUUoOeZ9t2YftQz4NgldqhnjtdTeAoAM2BtoNkXR9Xu7l6AKkDEEIIMXY33HADX/va1wr3810Ao33ubbfdVggE070LYEKHAXqhE+w/k0Rrjeu70goghBDignz0ox+lo6ODDRs2sHnzZu67775RP/ed73wnixYtYtOmTWzevJkf/ehH43ikE2/ClgOuWF6h6973Il664O9569XzKLPLqY6UUxmzJ+R4hBBCDE+WA55Y02o54IgZJVR2iidPdgPBhEAZz5uowxFCCCFmlAkMAGF8o5uj7afpSbu4voPW4Mi0wEIIIcS4m8AAEAFARU6ypyGBp3087UkdgBBCCHERjCoAKKVuUko9q5Q6rJT60DD7bFdK7VZKPaOUeuBcrxk2bMJmhFDZKXaf7AVkXQAhhBDiYjnnPABKKRP4OvBCoAHYoZS6Q2u9r2ifKuAbwE1a65NKqTmjefPlFSs44pziqVwAcHwXx/Nn3FhMIYQQ4mIbTQvA5cBhrfVRrXUW+DHw8gH7vAH4hdb6JIDWunk0b74qvpys0UhDZw8tPdmgDgBkeWAhhBBinI0mANQBp4ruN+S2FVsFVCul7ldK7VJKvWU0b746vgyNhxlpYPfJXlzfRWuN48mEQEIIIfrr7OzkG9/4xri/z+tf/3o2bdrEv/3bv437e02k0QSAodriB56hLWAr8GLgRuBWpdSqQS+k1LuUUjuVUjvb2ztYGV8GQFllQ6EOQKYFFkIIMZThAoBXwiHkZ8+e5U9/+hNPP/0073vf+0b1nPzMgaVSyp9nJKMJAA3AwqL79cDpIfb5rdY6obVuBR4ENg98Ia31t7TWl2mtL6upqaHCjjM/Op+KqiAABDMCOoU6ACGEECLvQx/6EEeOHGHLli1s27aN6667jje84Q2FxYFe8YpXsHXrVtavX8+3vvWtwvPKy8v5yEc+wubNm7nyyitpamoC4Pbbby/MGHjNNdcAwXTAzc3NbNmyhYceeojdu3dz5ZVXsmnTJl75ylcWphbevn07//iP/8i1117LV77yFbZv38773vc+rrnmGtauXcuOHTu45ZZbWLlyJR/96EcLx/KDH/yAyy+/nC1btvCXf/mXhZN9eXk5H/vYx7jiiit45JFHLsrnOZrFgHYAK5VSS4FG4HUEff7Ffgl8TSllASHgCmDEtpN8jd+qiuU8lnmazqTDibYMK+eEAGR1QCGEmKS+8PgXONB+oKSvuaZmDR+8/IMj7vP5z3+evXv3snv3bu6//35e/OIXs3fvXpYuXQrAd77zHWpqakilUmzbto1XvepV1NbWkkgkuPLKK/nsZz/LP/zDP/Dtb3+bj370o3zqU5/i7rvvpq6ujs7OTgDuuOMOXvKSl7B7924ANm3axFe/+lWuvfZaPvaxj/HJT36SL3/5y0DQIvHAA8Ggt1/96leEQiEefPBBvvKVr/Dyl7+cXbt2UVNTw/Lly3nf+95Hc3MzP/nJT/jjH/+Ibdu85z3v4Yc//CFvectbSCQSbNiwgU996lMl/VxHcs4WAK21C7wXuBvYD/xUa/2MUurdSql35/bZD/wWeBp4HPgPrfXekV63EADiy0n73Si7nd0ne/B00JQihYBCCCFGcvnllxdO/gD//u//XviWf+rUKQ4dOgQESwG/5CUvAWDr1q2F5X+vvvpq3va2t/Htb397yGb3rq4uOjs7ufbaawF461vfyoMPPlh4/LWvfW2//V/2spcBwXLD69evZ/78+YTDYZYtW8apU6e499572bVrF9u2bWPLli3ce++9HD16FADTNHnVq15Vok9mdEa1HLDW+i7grgHbbhtw/4vAF0f7xgqFAlZVrABg1qxGdp9cwisunY3nu2RdGQYohBCT0bm+qV8sZWVlhdv3338/99xzD4888gixWIzt27eTTqeB/ksBm6ZZ6LO/7bbbeOyxx7jzzjvZsmVL4Vv/WN4f6LfE8MBli103KHJ/61vfyuc+97lBrxWJRDDNi9vqPWEzASo0lmmwMFZH1IxSXXOavQ0JXE/j+C6eL8sDCyGE6BOPx+np6Rnysa6uLqqrq4nFYhw4cIBHH330nK935MgRrrjiCj71qU8xa9YsTp061e/xyspKqqureeihhwD4/ve/X2gNGIvrr7+en/3sZzQ3ByPl29vbOXHixJhf70KNqgVgPJhtR7ANcJTBivgyziSPkXJ8Dp5NcumiKBAh6/pEQ1IHIIQQAmpra7n66qvZsGED0WiUuXPnFh676aabuO2229i0aROrV6/myiuvPOfrfeADH+DQoUNorbn++uvZvHnzoBPy9773Pd797neTTCZZtmwZ3/3ud8d8/OvWreMzn/kMN9xwA77vY9s2X//611m8ePGYX/NCTNhywJctMPWdv7ud9thifnL8F/zi1K/oefYTvOHyRbz56gVUh2uI2CaVUVkeWAghJposBzyxptVywABlDY+hVFAHoNEsmt8cFAL6Hr72ZT4AIYQQYpxMXACwo4RPPIJtGKyMLwegdtZpDpxNksx6uL6DrzWujAYQQgghSm7iAkA4jnV2DyGvl3K7jLroAnz7GJ4PexsSOL4DyHBAIYQQYjxMXAAIxVHap6xxJxBMCNSUOUzIpLAuACDdAEIIMUnIDK0TY7w+9wkMAFF0pIroqUcxlGJ1xQp63V5W1CeCAKCDMZPSAiCEEBMvEonQ1tYmIeAi01rT1tZGJBIp+WtP2DBAUKglzyN8/CGsqzWr4sGEQHNmN3L/iTLaerNUhFxsbZN1fULWhNYrCiHEjFZfX09DQwMtLS0TfSgzTiQSob6+vuSvO4EBAFh6DcaBX1HecZAF8RWUWTFU+BiwiqdO9rKwqgrbsMl6EgCEEGIi2bbdb9pdMfVN7Fl1yXPRyqDs1GMYuQmBmtKHiEdMnjrVWygEdKQOQAghhCipiQ0A0Wr0vE2ETz6CaShWx1fSkGxk/UKD3Sd7cbwsgCwPLIQQQpTYhLerq6XXYrccIJztYFXFcjSaBbPP0NLj0NCRwdMeGhkOKIQQQpTSxAeAZcHCCvHTO1gRX45CYUaOAPDkyR7c/HwA0g0ghBBClMyEBwBmr8Uvm0Xs1KOUWVEWxuo4kz7MnAo76AaQ+QCEEEKIkpvAAKByVwp/yfMIndqBqXxWVazgUPdhNtdHePpULxk3qANwfY0vywMLIYQQJTFxAUCpvptLr8XI9lLe+gyrKpaT9FIsmdtOIuNz4GxPoQBQ6gCEEEKI0pjAAND31ubi56INi/KGxwoTAtmRowC5boCgDiAj3QBCCCFESUxsDYBpB9fhcrwFW4ieepQF0XmUW2U0pA6zdFaE3VIIKIQQQpTcBAeAUOGmv+R5WO1HiaVbWVWxgoPdR9iyMMK+00l6MplgH63xpA5ACCGEuGATGwCscOGmWrYdgPjpx1gVX05j6jRr5mZxPc3TjV2F/aQVQAghhLhwE98CkCsGtGpW4sUXEGt4jFUVKwEIR45hGYonT/TI8sBCCCFECU1sAFCq0A2gDANnyXMIN+5iVVk9CsXx5BHWzo/0LwT0vIk8YiGEEGJamPiJgKy+NY79Jdeg3DSz255lcdlCDvYcZnNdmKPNKdp6kwBoHawNIIQQQoixmwQBoK8OwFh0BdoMBXUAFSs41HOULfMMNLDrZGdhP+kGEEIIIS7MxAcAwwTDAsCyy8nWbSV68lFWxZeT9tJEy84SDRk8caIbTwfN/xIAhBBCiAsz8QEACq0AlmHiLHkOVncjG81yAI70HmFzXZTdJ3sL8wHI8sBCCCHEhZlUAQDAX/JcAFa2HqLCjvNs92G21Nmc7cpysr0XAA04ngQAIYQQYqwmRwAoGg5oVi7GrV5CvPHxXB3AYS6ZH3QR7DreWXiKrAsghBBCjN3kCABFwwFtwyKz6CrCZ3azOraYM6kmKsuS1JaZ7DrZ1bcwkNQBCCGEEGM2OQIAFLoBbMMmu/gqlO9yiRMU/R3qOcKW+mBdgKzUAQghhBAXbBIFgGA+AKUU/oIt+HaMS9uOYSqTgz1HuGSBTXfK49mzfdMCy+qAQgghxNhMngBQNBzQtmJkF15OTeMOlpQt5FD3YbYsCFYO3HG8vfAUqQMQQgghxmbyBAAAq68OILvoKsxEC2tDszncc5SqGCysttl1XBYGEkIIIS7UJAsAQTeAbdhkF10JwCWZDBk/y4nEKS6pC7G3sZdUNguA58vywEIIIcRYTK4AkBsOaBkmlM/GmbWaba3HADjYfZhLFthkXM3uho7CU2RdACGEEOL8Ta4AUDQc0FIW2cVXsrhpP9V2BQe7j7BxQQhD9a8DkEJAIYQQ4vxNrgAARcMBLTKLnoOhfdZZVRzsOUxZyGD1nBC7TnQWdpc6ACGEEOL8Tb4AYObXBbBw56zFj1RySTpFc7qFzmwXW+pCHDybpDsV1AH4WuNKN4AQQghxXiZhALDAsLANCwyT7MIruKylqA6gLoSvYceJ1sJTZDigEEIIcX4mXwAAsEIYysA0DDKLrmJDbzsWBgd7jrBmjk3YUjx+rGg+AOkGEEIIIc7L5AwA+W4AZZFddAUhDFaaZRzsPoxtKjbOD7HrRNF8ANICIIQQQpyXyRkArDAohW1Y6Eglztz1XJJKcqT3GK7vckldiFPtaZq60wBoLa0AQgghxPmYnAEgNxzQNoLpf7OLruLSjrM4vsPxxEm21AVDBR871lcHIPMBCCGEEKM3OQMA5AKAhVKQXXwVmzNB1f/B7sMsqbGojBhSByCEEEKM0eQNALlpgS1l4dauZHa4irna5GD3YQyl2FIXYteJjsKSwLI8sBBCCDF6kzcAmMEwQNuwQCmyi65icyrJwZ7DAFxSF6I94XKsNQGARooBhRBCiNEaVQBQSt2klHpWKXVYKfWhIR7frpTqUkrtzl0+VpKjs8KFOoDM4uewJZWkNdNOe6aDLXXBSIHHjrUVdpduACGEEGJ0zhkAlFIm8HXgZmAd8Hql1Lohdn1Ia70ld/lUSY7ODActAIBTfxmbsi4Q1AHMjZssqDB5XAKAEEIIcd5G0wJwOXBYa31Ua50Ffgy8fHwPK8cKYxgmhjLQoTJW1K4lpCl0A2ypC7H7VFdhKmDX1/iyPLAQQghxTqMJAHXAqaL7DbltA12llHpKKfUbpdT6khxdYThg0ArgL76a9Zk0hzr3A3BJXZhk1mffme7CU6QOQAghhDi30QQANcS2gV+znwAWa603A18F/nfIF1LqXUqpnUqpnS0tLaM7wqIAkF10FZvTWY4mGnB8h80LQijoPxxQAoAQQghxTqMJAA3AwqL79cDp4h201t1a697c7bsAWyk1a+ALaa2/pbW+TGt92ezZs0d3hFakEAC8qkVsMspw8DnWe4J4xGDFLIvHiyYEkjoAIYQQ4txGEwB2ACuVUkuVUiHgdcAdxTsopeYppVTu9uW5120b9EpjYVpYZhilAKVYM3crAAe7ngVgS12Yvad7SOYKBD1f40kdgBBCCDGicwYArbULvBe4G9gP/FRr/YxS6t1KqXfndns1sFcp9RTw78DrdAln5VFWGEsFrQDxJddQ57gcaXkCCOYD8Hx48mRnYX9pBRBCCCFGZo1mp1yz/l0Dtt1WdPtrwNdKe2hFrAiWYeH4Ltm6S9m002FH4iRaa9bNC2GbsPN4O1evCHodsq5PNGSO2+EIIYQQU93knQmwmBXGNu3C7Q3R+bTi0JZuI2wp1s8L9asDyHieTAsshBBCjGBqBAClsK1o4e7aeZcBcLT5cSDoBjjckqKtNwMEywP3ZNyLf5xCCCHEFDE1AgBg2jEMFRzuwmU3EvF9juQCQH5a4J3H+4YDprIeGde7+AcqhBBCTAFTJgAE6wIE/fpGVT3rfYMDyQYAltdaxMOKHcf6DzzoTrnSFSCEEEIMYeoEANPGzi0RDLA+Vsch5ZDNdGEaik0LQjx+vL3fCd/Xmu60dAUIIYQQA02dAADYVqxwe838bbhKcebE74FgWuCmHodTHal+z0k7HmlHugKEEEKIYlMqAFihssLt1UteAMCRlp1AsDAQwI6jrYOe1512ZJEgIYQQosiUCgDKimDl6gCqorUsxOLZZANozYIKkznl5qA6AAhGBXSnnYt9uEIIIcSkNaUCAIaBbfd1A6wvW8jTliLUfgSlFFvqQuw62TXkVMAZ15euACGEECJnagUAwLb7ugHWzNtGm2XSc+p+IJgPoDvjcfBs15DP7U47sk6AEEIIwRQMAMV1AOvmXgrA0ZZdAGxeENQBPH506KWGtYbulHQFCCGEEFMvANhRjFwdwJL4QmKY7M+2YGS6qY6ZLKu12XG8Y9jnZz2/sHKgEEIIMVNNuQAAfd0ApjJYU76Qp8Ihyk/3jQZ4qrGXdHb4b/q9aRfXkxUDhRBCzFxTMgBYob5CwHVzNnMwZGM0PALAlgUhsp5mz9HGoM1/CBpkgiAhhBAz2pQMALZdXri9rno1nlKcaH0SfI8N820sAx4/0YnKdA/7Go7nk5AFg4QQQsxQUzMAWGHILQ+8tnIlAHsNl0jbQaK2wWULo/z86Q4a27rASQz7OomMiyNdAUIIIWagKRkAlFKYdrAuQEWonIXRuTwZCVPe+BgAf/3cCpSCj93dgJ/qBjcz5OtoglEBsmCQEEKImWZKBgDo3w2wvmYNT0djlDUGdQCzyhT/cN1c9pxN8R+Pt2BkusAfurnf9TW90hUghBBihpm6ASAUA6UAWFe1kk6laek6hplqB+CKJYqXrq3iuztaeKKhB5XuAD10c38y65F1pStACCHEzDF1A4BhgxVM/LOuahUAT0XClDc+DoCvff7q6grqK0PcencDPcksKt057Ot1p6UrQAghxMwxZQOAZVgYVlAHsLi8jjIryhNlFZQ3PlrYRxlZPv7CebQlXf7pD6fBzaAyPUO+nudreqQrQAghxAwxZQMAgJWbEMhQBmsrV/JUWZyy07v69fcvrPX4qyvncO/hbu7Y14lyEuAkh3y9VNYj48qCQUIIIaa/KR0AguGAFhDUARzVGVJukljzM4V9PN/jls0xti0s418eOMPxjgxGpge87JCv2Z1y8WXBICGEENPclA4AlmEV1QGsxEezJxrt1w0AkPHS3Hr9PMKWwUd/20DW9TDSneAP/rbva02PzBIohBBimpvSAcA2bDDCAKytWoFCsat2EWW5+QCKRSMZbr1+Ac+2pPnmo82g/dzIgMHf9tOuR9qRrgAhhBDT15QOAIYyggmBlKLcLmNxeR1PxcqJdB7H6j3bb1/Xd9m22OLVG6v5wRNtPHqiF+W7w44M6E470hUghBBi2prSAQDANu1+3QDPeD34UJgVsFjKS/Peq2eztCbMx3/fSHvSRXkZVLZ30L5aByFACCGEmI6mfACwlAVGEADWV62ix0txqGoB1QfvRA0o9NNa45LkszfV05vx+PQ9jWitgwDgpAa9dsb1SWWlK0AIIcT0M+UDQNACENQBrKsKFgZ6bPXziXQcZs7Obw7a3/Fd6qs0/99z5/Lw8V5ufzqYOdDIdIM3+Bt/T9rBk64AIYQQ08zUDwCGjTJMMC3qy+YTt8vYYxu0rXsNNc/+koqj9w56TtJL8aoNlTx3STlfebiJw61pQGMMMV1wfsEgIYQQYjqZ8gEAcsMBTbswIdD+rkN0XfEuknM2Mv/RLxHqPN5vf601CS/Bx15QRzxs8pHfNpB2/WFHBmQ9n2RWhgYKIYSYPqZFALANG8xgWuD11Ss50duINj3OXvdxfCtK/QOfRA3o43d8l0jY5RMvrONoe4Z/f7gJAOU5qEz3oPfoTbu4niwYJIQQYnqYHgHAtMGyQanCwkAHOo8QrZ7H6Ws+Qqj7FPMf+dKgb/ZJN8m2RVHeeEkttz/dzgNHgxO/clOobKLfvhrokq4AIYQQ08T0CACGHdywQqypXI6BYl/nIUylUEsup2Xz26g8/geqDt7R73laa3qdBO+5ag6rZ0f49D2naekNTvIq2wNuut/+rq/plQWDhBBCTAPTIgAYysBQBhghYlaUFRVLuPfMH8l4WcKmSeLSN9Fbdzlzd3yDSOuBfs91fAdfZfnMTfVkXJ+P/74RP9dSYKS7+i0sBJDIuDjSFSCEEGKKmxYBAHKtALkJgd656vU0Js/yX4dvB6AsHKbpmo/gRaupe+CTwZC/Ikk3xaIqm7+/Zj47TiX44RNtuUd0riiw/wm/M+mQdSUECCGEmLqmVwDIDQfcOmsjL66/np8du5N9nYcAiFXW0njtx7FTbSx4+PP9Tuq+9ul1E7x8fRXPX17B1x9pYl9TUDSofG/QdMG+1nQks/RmXPQQawkIIYQQk930CgAAZnD9l2vewKxIDV/ccxtZL4upFNbCjTRd9h7ijY9Su/fH/Z6f9bJk/Sz/eP18amMWt97dQDI3C6DyskOODEhkXNoTWRkdIIQQYsqZNgHAMiwUCsxgVsAyK8bfrf8LTiZO89+Hfw5AyDBJb7yFriXXMXv3d4idfbLfayTcJPGwwadvrOdUZ5YvPdi3oJBykuD0HxkAQWFgeyJLQooDhRBCTCHTJgAopQoTAqGCbdtmb+bm+uv4ybFfcaDrCACxkE3rcz9AtqKeugc/g5VsLbyGr30SbpJL68p4+7ZZ3LGvk98f7Co8bmR6wc0Mem8N9OZaA2TaYCGEEFPBtAkAkJsRUCkwQ4Vt7179Jmoj1UFXgB8M8Ssrr+D0tZ/AcFPUPfhp8PsW/Ml4GbJelr+4fA4b50X5pz+c5kx3flEhjZHp6rd/McfzaevNyAJCQgghJr1pFQD66gDChW3ldoz3rf8Ljvc28IPDvwDAUIrQ/JWcufLviDXvYfaT/9nvdXrdBIah+dSN9WgNt97diJv/Zq99jFRbUBPgJIMFhIoKATXBMsKdySy+tAYIIYSYpKZnALBC/bZfMXsLN9Zdy/8cu4ODXUdz+xq4a26kY9VLmfXMjyk/+cfC/r72SbpJ6itDfPC6+Tx1Jsl3d7T0vaD2UU4SI9ONkWrDSDSjkq2odFdQJ+BlyTgerYkMaUdaA4QQQkw+0yoAmIaZmxDIhHCs32N/teZNVIcq+Oc9t+HkJveJ2hbtV7yXVO0qFvzx89g9pwv7p70MWd/h5jVV3Ly6kv94vIWnTieHeWeN8l2Um8LI9GCk2jESTZBoobujha7ODvxsGnwZLSCEEGJymFYBAIpaAULlhSGBAHG7nL9d/06O9Z7ih0f+X2F7WVkZZ679BChF/QOfQHnZwmO9Ti9aa/5h+3zmxW1uvbuB3szov9Er30O5abLJLjramsh0nYbeZki2Q6YnKCiUUCCEEGICTN8AABCpCIoCc54zZysvXPA8fnT0lxzuPg4E9QDh2Qs5c/WHiLQfZu7jXyvsnx8VUB42+cyN9TT3OnzuvtNjmvzH15rulEtPKo1205DpDYJAb1P/UOBIS4EQQojxN70DgGFCON7v8feseQuVoTj/vOc23FxXgG0YeMufR+uG11N96NdUHvldYf+0l8bxHTbOj/GXV87hdwe7ufNAF2OVdnw6kg7Z4smDfC9oDcj0QqoDEs2Q7h60eqEQQghRKqMKAEqpm5RSzyqlDiulPjTCftuUUp5S6tWlO8TzYxlW/w12BOy+UQEVoXL+dt07ONJzgh8d/WVhe9S26Nr6DhJzNzPv0X8j3HG08Fivm0BrzVu2zuLSuhj/fP8ZDrb0XynwfHi+pivlkMi6aIY4yWsN2QQkWsBJjfl9hBBCiOGcMwAopUzg68DNwDrg9UqpdcPs9wXg7lIf5PkoTAhULFwBRt+PevXcy7h+/tX84Mj/40jPicL2smiEM9d+DD9URt0Dn8BwgqI/z/dIeilMQ/HJG+qJWoq3/eQo393RguuN/Vt6MuvRkXRwh2vy9z1IdUKiLRhuKIQQQpTIaFoALgcOa62Paq2zwI+Blw+x398APweaS3h8Y9KvGwCCOoBwZb9Nf732rcTtMr645/8WugIMpYhVz6HxebcS6jnN/D/9S6EZPuWmcH2XeXGbH71xBdcui/ONR5p520+PXnBrQGfSIZkdYSphLwuJ1iAMSH2AEEKIEhhNAKgDThXdb8htK1BK1QGvBG4r3aGN3aAAAGDZ/YYGVobi/O26d3Co+xg/Ofarvt0MA7X4UloueQcVJ+6n+kDfiIF8V0BtzOJzL1rIF160kJZeh7f85Aj/99FmnDEuCqSBRNajI+XgjdTv76SCboHs4DUJhBBCiPMxmgCghtg28Cz1ZeCDWusRx8gppd6llNqplNrZ0tIy0q4XZFAXQN6AoYHPm3c52+ddxX8f/jnHevoyTsSy6Nn0enrqr2LurtuItOwDwPVdUl5fn/zzV1Twkzet4IaVwTwBb/7xUZ5pGnufvev5dCSypJxhagMgWMY43Q29LUOuSyCEEEKMxmgCQAOwsOh+PXB6wD6XAT9WSh0HXg18Qyn1ioEvpLX+ltb6Mq31ZbNnzx7bEY+CZViEi6YD7mfA0MC/Wfc2yu0yvrjnNryiOf7LoyGanvdhnNgs6h/4FGY6qPxPeel++1VFLT51Yz3/+tJFdKc9/vynR/nqH5vIuGNvDejNeLQlsvRkRqoPcIOhg8n2YdcmEEIIIYYzmgCwA1iplFqqlAoBrwPuKN5Ba71Ua71Ea70E+BnwHq31/5b6YM9HPBQPZgUcaMDQwKpQBf/furfzbPdRfnr814XtCkWsoprGaz+Bme5gwcP/BNpHa02P24uv+5+Yn7c0zk/etIKXrq3iv3e18sb/OcJTZ4abOfDctO4bMtiZcki7w5zk3UzQLZDpkWGDQgghRu2cAUBr7QLvJaju3w/8VGv9jFLq3Uqpd4/3AY6VoQziofjQDw4YGnjtvCu5Zu4VfO/Qzzje21DYbhkG5oJ1NF3+XspP72DWnh8CQVdAl9ONN6DHIx42+egL6vjqyxeTcTV/cfsx/vXBM6ScCyvcczyfnrRLWyJLMuviDzzRax3MISDDBoUQQoySGsusdqVw2WWX6Z07d477+3Rlush4Q/SVaw3JtkJVfUemiz9/+P0siM3l36/4JKZhFnbtSWWZ9eBnqTh6Lydf8M8kF2wFgpBRYceHrDlIZD2+9scmfrang7oKm1tfUMfW+rKS/EwKCFkG0ZCJbQyR4cxQ0NVhDlEMKYQQYlpQSu3SWl821udPu5kABxq2K2DA0MDqcCV/s+7tHOg6ws9O3NVv1/KoTfNVf0emajF1D30GKxEUMPrap8vpJusPHqNfFjL54HULuO2WJSilePcvjvP5+06TyF54f70GMq5PZ9KhI9c90K9oMD9sMN0lwwaFEEIMadoHgBG7AgYMDbxu3lVcPecyvnvodk72Nha2KxRl5XFOX/sJDC9D3YOfCorwIKgJcHqGbmUAttaX8T9vWM4bttTyiz0dvO6HR3j0RG/Jfj63qHsgkXX7DyPMJmXYoBBCiCFN+wAAEDbDRKzI0A8WDQ1USvG3699BxAzxxb3/F6+o0M8yDKw5yzhz1fuJtTzD8v99KzV7f4yZ7sqFgF5S7tD97xHb4H3XzOM//mwpEUvxN788wafvaaTnPFYWPBetg5kF2xNZulNFaw3khw0mWmXYoBBCiIIZEQAA4vYwXQHQb2hgTbiK9659G/s6D/GL47/pv5tlkl75Ak5t/xRO2WzmPvEtVvzsNcx/+HNEW54h4STodYb/tr1pfowfvH45b906izv3d/LaHxzmoWM9JfsZ8zKeT1fKoT1ZNKeA5wRDBlMd0i0ghBBi+hcBFst6WToznUM/6KSDb8oEzfq3PvEv7Grbw7ev/gL1ZfMLu2mChXw8XxPuOEbVwTuoPPp7TCdJunoFHatfRnrFzZRFZ6PUUHMoBfY1pfj0PY0cbstw0+pK/v6aeVRFh5nA6AIpBWHLIGqbWIYBhgWx2n7rIwghhJhaLrQIcEYFAIDubDdpd5i5+9PdQRAAWtPtvOPhf2BJeT3/esXHMItaDzytSWTcwtS/hpOk4ti9VD/7SyIdR/HsMnqW34i74Q3ompXDHovj+Xx3Ryvf2dlCZTgoGnz+iorS/bBDsE2DiG0QDoVRZbP6TYokhBBi6pAAcJ601rSl2wZN5JN7sN/Mer9rfJAv7Pkm71nzFl615OZBu2c8j0TG7Zt/R2uiLc9QffAO4scfwPAdsvO3kl73GjJLrh92WN7BljSfvqeRAy1prl9RwQe2z6c2Nj6tAXlKQTgcIVIxG9syz/0EIYQQk4oEgDEYsSvAdYJ+coKw8JEn/pndbfv49tVfoK5s3qDdfa1JZF2yA6b+NdOdVB7+LdUHf0Wo9wx+tJbUmltIr30Vfvn8Qa/j+prv72rl24+1ELUNbtlYzS0bqplfEbrgn3ck2gxhxGqIhiwitolpSIuAEEJMBRIAxqgn2zNs1T7ZXsgE0/i2pNt5x8MfYHl8EV+6/NZhCwkd3yeRcfH8gbP0+cRP72LWoV8TPvkwKEV20TWk1r0Gp/4qGPB6x9ozfP1PTYXiwKuXxPmzTTVcsagMY5ya67UZRkerAQiZwQRDYcsYsYZBCCHExJIAMEYjdgUAJDuCynngtw3388W9/5e/Wft2XrH4huFfE03K8UhnvUFr+SmlqEr1ED94B9Fnf4GRaserWEhq7atJr345OlLdb/+zPVl+saeD/32mg46UR31liFdtrOal66qojJS+e0BbEXSkqu94gbBtBvUC0kUghBCTjgSAC+B4Dh2ZjqEf9L2gHkBrtNZ8eNfn2dPxLN+++gssiM0d8XVd3yeZ9QpFgsXK7DKimISP3UNk308JnX0CbYbILLuB1LrX4s7Z2K8wL+v63Hekh5/taWf36SRhU3HD6kpevbGGdXOjF/TzD6TtGDo8uAjRUIqInRtFYMrIASGEmAwkAFygEbsCioYGNqfaeMfDH2BV5VK+uO0jw88pUCTteiSz7qBF+qJWlDIrmIHQbD9EdN/thA/9GsNJ4NSuJr3utaRXvWxQ0eDBljQ/39POb57tIuX4rJ8b5dWbanjBygoiVmlOzNouQ4eHmTmRvlEEEcvEkHoBIYSYMBIALpDWmvZ0+6CV/QqKhgbeeeoP/Osz3+Zv1/05L130wlG9vp8bMpgd0BoQMcOU2+WF+yqbIHz4LqL7forVfhC3chGJK/+e7KJrBw3V68143Hmgk5893c7xjiyVEZOXravilo011FdeeNGgDpWjQ+Uj7pNfkChiS72AEEJMBAkAJTBiV0DR0ECtNR/c+Tn2dR7iW1d//pxdAcWyvkcy4/UrEgyZIeJWef+Tp9aETj1E2SNfwuo6TrbuSnqvej/eEPMJaK3Z2ZDgZ0938MDRbnwNVy0u58821XDV4vILquj3w3GwR7d6oVIQsU2itoktXQRCCHFRSAAokd5sL0k3OfSDRUMDm1It/MUfP0TIsPnkJX/H+upVo36PoYoELcOiYqhpin2H6L6fEtt1GyrbS3rtq0ls/St0tGbI127udfh/ezv4f3s7aEu6LKiwuWVDNS9fXz3mGQb9cAXYsXPvWMQ0FGHLIGQZhExpGRBCiPEiAaBEztkVUDQ08HhvA7c+8S+0pNr4uw1/wQ1115zXe7m+TyLr4ea6BUzDpMKOY6rB1fYq3UnZrm8S2Xc72o6RvPRdpNa/fthJhVxPc//Rbm5/up0nGpOETMULVlbw6o01bJgXPe8Tsh+uBHtsxYaKoGYglAsE0joghBClIwGghEbsCoB+QwO7sj18aveX2d2+j9csfQnvXPX6ftMFj0ZxkaChDCrsOJYx9Ld1s+MI5Y/8C6GGP41YH1DsSFuan+/p4K79nSQcnzWzI7x6Uw03rqokYo/2WBV+pBKGW03xPCgFYdMsBAKZdEgIIcZOAkCJjdgVUDQ0EMD1Xb6+/3vcceoerph9CR/Z/N5Cdf9oFRcJKqWI23FCxtDf7gFCJx+m7NF/weo8NmJ9QLFE1uM3B7q4/el2jrZniNkGz11azvOXV/CcJXGi5wwDCj9SBVb4vH62czENVegqkEJCIYQ4PxIASuycXQFFQwPzfnnyd3xt//dYWDafz1z6gfMqDszL+h6JdFAbEDOjhM3w8EMNfYfIvtsp2/XNoD5gzatIXPaeYesDin+2J08nuetAJw8c6aEz7RG2FFctDsLA85bGKQ8PN+mPwo9Wgzk+UxMrwDL7agdCJRrWKIQQ05UEgHHg+A4d6RG6AoqGBuY92baXT+7+MqD4xJb3saV23Xm/r0aTzLqknaA1wDZsIkYY27CH/Has0p3Edt1GdN9Pc/UBf0Fq/RuGrQ8o5vqa3Y1J/nCkm/uOdNOacLENxRWLyrhuRQXXLI0PLh5UBn6kelSvf6GUohAEQqYhExAJIcQAEgDGyYhdAQDZRHAp+vgaE2f56BNfpDHZxN+sfRsvXfSCMb236/tkXI+M6xfqA0JGiLAZwh6ie8DsOErZo18ifOph3IpcfcDikesDivlas+dsij8c7ua+w92c6XEwFWytL+P5KyrYvryib3VCZeBHa2CYWoXxYqiguyDoKujbrui703970e0Bn0P/x4bfTwghJjMJAONEa01HpgPXd4ffyfcg0wNutrCp10ny2ae+yuOtu3nFoht4z5q3YBpjm0tfo8m4PhnX7zdiIGyECBvhQa/bvz7gCnqv+sA56wMGvafW7G9OB2HgSDcnO7MoYMuCGNetqOC65RXMqwjjR2thjD/XZFbcFWGbSoYyCiEmLQkA4+icXQGFHdNBEMh9lp72+fazP+L243dyae0Gbt38f6g4x8x65+L6PinHw3H9QqODbViEzDBhI9RXL+A7RPb9LFcf0EN6zS25+oDa835PrTVH2jL84XA3fzjSzZG2DADr50Z5/spKrtuwhLqa0U0WNJVZhsLOdUXYpoxeEEJMDhIAxlnCSZBwEufe0feDuQKKagN+23A///bMfzA3OpvPXPp+FpXXXfDx+FqT9TzSjl+YVTBfLxA2QoSMEEopVLqL2BO3EX3mJ2grSnLru0ZdHzCcEx2ZQsvA/ubg51w1t5zrVs/hujVzWDpr+ocByHVHFLUSSH2CEGIiSAAYZ6PqCijmOpDpDroHgL0dz/LxJ/8Vx3f56Oa/4fLZW0p2bI7vk86tOpj/rxjUC9iEzDAhw87VB/wr4VMP5eoD/o7s4u2jrg8YzunuLPcd7uYPR3p4+kxQK7GkNsZ1a+ZwzcrZrJ4XnzHflPMFi3bhoqTbQAgx7iQAXASO79CZ7kQzys9Ka3ASkE2CDqYPvvWJL3Gs5yTvXvMmbll8c0lPEL7WZNygVcAv+u9pKIOwGSZihIk0Pkr5I/+C1XmUbP1z6L7us+ccNjhazSm474TDfc+28OTJDnwNFRGLrYuruXxpDZcvraG++vzmR5jKpI5ACHExSAC4SEbdFVDMc4PaAM8h5ab5/J5v8HDTDm6q287/Wf/nI074M1YZzyPr+INWH7QMizAGVc/+kvjjX8GPzqLrxi/j1a4uyftqM4SOVNORdHj8eDuPHwsuzT1B3cCCqgiXLwnCwGWLa6iMjf9QwskkX0dgG8EoBqWCEQyGCrpw8teTmu+DId0dQkwWEgAuovZ0++i7AoplU5Dtxfc9vnf45/zgyC/YUL2aT2x5H9XhytIfKOBpTdpxC0MJi1W0H2XeHz6Mke2he/tnyC4b3dLG56KtCDpS1Xdfa062J4MwcLydXSc6SGQ8FLB6XjxoHVhSw6aFlYSt6Tei4Hyp3P8ZKhjcaCiVCwvB9eDtQ+870Gj+jQ+7i9bgpsFJgZfBsMIQjoM1PhNCCSFGTwLAReT6Lh3pjtF3BRTzfcj2gJPhD2f+xBf33EZ1qJJPb30/y+OLS3+wORpN1vNJO31DCQGsVDsL7/8EkZa9JC59F8mtfwXnuZbBkO9nRdGh8iGHCLq+z/7TPTx2rI0dxzvY09iF52vClsHmhVWFQLBybjnGZP82PN15DspNodw06L7fG6XAMgzMUBg7WoFlh6UIUogJIgHgIks6SXqd3rG/gJuFTDfPdhzi1ie+RMJN8uFNf81z524r3UEO99a+T9r1yOZaBZSXZd5j/07V4bvILN5Oz3WfDU7epaIMtGGCCi7aMHK3g+uEo3nyVCePH2tnx7F2jrYGXSxVUZvLllRzxdJati2tZn7l2FYjFOdJ++CmUE4KNcqWLm2GIVyOZYexTYVtGliGjIwQF4HWweJsXjb4wmFFLri4eaqRADABOtIdOL4z9hfQGrK9tHY38rEn/5Vnu47w5ytfyxuWvfyi9QM7vo/j+TiuT/yZnzN35zdwKxfTfdNX8SsWXpRjgKAdOx8SWhIej5/q4fET3ew40UVrbzDB0sKaaKF+YOviauKRmVU/MO7cTO7bfgbG0rpFEAR0OF6YITJfCGmZCtuQ4ZKiBHw/ONl72eDE7zv9+66UCkKAHS35wmWTlQSACeD5Hl3ZrrHVA/R7IZdMsoUvPfVN7j3zR54//zm8f8NfEh6nBXeG42uNcWoHtfd8HLSm4/mfQy963kU9hoG01hxtz/D4qQSPnUrwREOClBM0Rdum6uvzZnAfuSpcD+hHZ3D/ORQ/D8KWyep5cTbVV7KpvpJFNbHJX5w3Fr4XnPSdZL8m/gulrUiuC2jwVNEDQ4GVazG44PfUGq2D6OIXbuvCuUFr+o2Ogb4visNNJT3UtqH2LX6Kacjwz5LKf7svnPCHWaBtKMoAOwJWdFrXq0gAmCBaa3qcHtJu+tw7n+u1Mgn+58AP+c+DP2F15XI+dcnfMStSmiF658PobqTqNx/C7DhOzxXvJbvp7YAxxu+EpeV4PnvPpniiMUnS8YM/+oCvAQ0+oFH4AFrho3KPK7RSuZNC/vHciYK+k0f+BNGbcdl3upvudBDuqqI2G3NhYFN9FWvnx0cuWNR+7luJ7rutc0eXew9tmMEJUpkXt8kyV9Cn3BTKy557/wt5qxFqQYoVhwLLUIUTef6/r/b7TubF2/3chsnwu1nMUMHPYeZ+HtNQQc3EDJkTY8y07n+y97IjVKaeJ8PqCwPmxV3DZLxJAJhgKTdFb7Z3bIWBxXyPP564l3968svErCivXvIiblhwzbiNEhiOcpLE7/00kWMPklp9M9nrPoJlV+D7iqyrB32TmnKUgVaqXy1CsM0ITlbKwPc9TrQlebqhi6cau9nT2MPJjiDoWYZizdwYm+bH2LygjI3zo8yOWX0n/bEcj2GBYeW6QqzgOEq52NIwBX3jT6HtKNoum5brRpwPBX1hoF84mKGtBr434IR/AV2q58O0+7oJpsHvpASAScD1XboyXXj6PJqohnGk9QD/vvur7O04gKVMrp57GS+uv55Latf3zfc/3rRPbOd/Ub7zP3HmrKP75n8iUrGImBXF0xrH88m4GrdoBsLpriPp8vTZJE+fSfL0mRT7mlJkveCnr6uwc4Egxub5MZbWhEvwjU/liiaDQNAXDqzR/eEaQ0Hf+FFoO4YOlZVkpMl0U9xqYCqVaw2Z+q0GWmscT+P6Pp6TGfDt3h+w71AvMKpNQBCwDKOvG8Y0gm5Cc6RwZYWDMGBFpuz8FhIAJgmtNd3ZbjJephQvxvG2/dx17Df8rvFBepxe5kfn8KL667ipfjs14aoLf49RCB19gIp7P40Oxei66XP48zYQs6LErKAqPz/E0HF9sp4urE0wE2RdnwMtaZ4+k+SpM0mePp2kPRUEwLKQwcZ5MTbNj7J5QYz1c6OUhUr5baM4HFh9Iy0MC3z3ggv6xo9Ch8rQdkyCwCgEJ7XBXQmWoTAmWTgoPtk7XvDlwPV18PuY6R73Lqfh5GuCDENh0j8kGLmgYBoGmKFc8eDUGkkgAWCSSTpJEk7iwrsEAHyPbKqDhxoe4s6GP/BU+z5MZfKcOVt5Uf11bJ21CXOc/5CabUeo+s2HMBIt9Fz7D6TXvAjTMCi3ygYVK3pak3VzgcAfPAHRdKa1prHL6QsEZ5IcacugAUPBylkRNs2PsXp2hKqISWXUpDJiURkxqYiYWJPsD/q4UkbQImCXTak/tpNJYT4Go393wsUYaeH7Gsf3cT2N6wW3B4V/rVFOApVNMPmC6GBGLiQYClQohmlHMUORQkjI/5YWij8nye+tBIBJyPEdujJd+KXqb/VcyPZyqvM4dzXcx92ND9Dl9DA3Moubc60Cs8exaFClu6j83a2EGneR3PQaeq/6azAsbMOi3I5hDzOlseP3tQ7MpO6CvJ6Mx96zKZ46k+Sp00meaUoVRjIMVB4yqIyYVEaDUDDwUjXE9qg9xdcYUAbazrcITOGfYxIpdavBqE72A3lZVKYLdT5V+5ORMoJRLVYkaCEYbrcBN/JxYeBIkfy/1YFhIhoyxzwTqgSAScrXPj3ZntJ0CeQ5Gcj2knXT/KlpJ3c2/IEn2vZioLhyzqW8uP75bJu9ZXxaBXyX8j99jdie28nWX0bXCz+NjlQAEDZDlFkxrBH6pmdyd0Ge62uaex260h5dKS+4Tru5a4/Owva+bYns8CHSNlSuJaHvErIG/Lcfsh918MaBfwaG+q+jgEXVYdbNibJuboRZZSWaj0EZ+KEysCQIjKfhWg3yTeK+H/wbdf0gsDveeRb9ah+V6UG5qfH7ISaINkywomir9MWDFRGb6Bi7CCUATHIl7RLIyyYhmwCtOZ1s4q5Tf+C3jQ/Qke1idqSGm+uCVoG50Vmle8+cyIE7iT/wRfzy2XTe/AW8mmVA8MclYkYos6KjKlb0dfDHJuuO4Q/NDOJ6mq6MR1eqLxT0XdyiIBEECNcb/DmO9pw6cL+BT3P9YBnofHabXWaxbm6UtXOirJ0TYd3cKFXRCxm9UDQxFEZulIZRNHLDGHwRJaHUBY66c1IY2Z6LPMpkYmjTDoKAFS1JYJUAMM05nkNXtoRdAlCYTRAnBToYifBI8xPc2XAvO1v3oIBtszfz4vrruXL2JZglTK3W2b1U3v2PKCdJ9/UfJ7u0b9KgoPI2eK+BTWH5U0r/CVWC236ufiDfzIgunoGlr+lMEby+qcyp3fw9RaUdn2db0uxrTrG/KcW+5hQnOvoKvBZU2KydE80Fgwhr50QpD4/jcKvicEDfUE6tFANDxLn/WF/g79NM/H30vaC5f4KK/CaWQlvhoIvAioz5VSQAzAC+9unOdJP1S/wPxfdyQaCvq+FssoXfNN7Hbxruoy3TSW24mpvqruVF9c9nXmx2Sd7W6G2h8u4PYzfvp3fbO0lufWtJv5G5+amKh6kfULmgYSkLy8hfT69JPqaK3ozHgeY0+5uDQLCvKcXp7r5x3YuqQqydG2VdrpVg9ewoUXt8v71nXJ+ejEdvxqc749GT8ci4OrfsMpi5GSEHXhsq6EM3oFAUZuRmjex/PXhbZcTCDlkDhm5O0KRPF4HKJlDZXqZCkd+4U0Yw+ZUVCeYaOA8SAGaQhJMg4SRK/8KeC5mefhNqeL7HY61P8utTf2BHy240cNmsjbx04Qu5as6lFz6vgJsh/sA/Ez34W9LLttPz/I8EBV0lptGFgkLXH35528GhwB6xLmHI9xpi2lhx/jpTLgeai1oKmlI0J4L5CAwFS2vCha6D9XOjrJgVIVxUv+D5mt6sR08mOJEXTuZpj56sR0862NaT9ftuZ4L9ezMemSG6QsaboaC+MsTSmjBLqsMsrQnnboeIhXItE0WTPWll9k36NJW6M7xsMLRvwueXmJy0YeXqBSKjqheQADDDZL0s3dnu0nYJ5OUKBQfOm92cauO3jfdxV8N9tKTbWVq+kLeseBXPnbvtwoKA1kSf/gnlj3wdt2ZZMF9AxYIL/CFKRykKrQO2YRXCwUjy9QnBWGZm5AiG8dCacNjfnGZfUypoLWhK0ZGbO8E0gpNnxtX0ZEYufoTgZBsPm7mLQXnYpCJsUh42iBduB4/l9wtbqjBttO9rfE3uogdcD96mdRBKNLlrHQx71bn9PR0UeB5vz3KsPcPJrgxFq28zL27nQkGoKBiEi2omVL8Wg+JJoCZN64H2UdneYP0IMSraDOXqBYafX0ACwAw0bl0CEPy1clKFQsFinva5/8wjfP/ILziVOM3S8oW8ecWreN4FBoHQyceo+P3HwTBIbPsLnNmrcWuWBpNrTDL5UGAbQTAYTSgIhkLluiT8mTmKodS01jT1OuxrCkLByc4MMdsgHjGJh4ITeEXEJB4y+m7nTuixST4E0vU0p7qCMHC8I8PR9gzHc7czbt/vTk3UZElNX2vB0powS6vDzCqz+v98yghOJqYNRui8m5kv2BiL/LrTHvubU5zpdoJ/Nzr4t+P5waiYYFvweQ18zPM1ru677fka16dw2yt6rCxksqDCZl7cZkFFiHkVNgviNrVlVmHRr4ml+oYUDlipUALADNab7SXpjlOi9n1wEoVCwWL5IPCDI7/gZCEI3MLz5l4+5iBgdp6i8rcfxuo4BgSL73iV9bi1y3OXFXg1y/AqFky6Jk9Dqb6ug1GEAmklEGPha82ZbofjHRmOtfddjndk6Mn0nVzLQ0bQSpALBPnb8+N2MEVwfg0JM4Q2Q2DY49NK4Hu5mfzOPZy5N+P1KxDd35ymoWt0X3CC4YjBtWUoTEVhiGIwXDGYwS8/VXK/xxV0ZzxOdwdDbIvZhmJe3C4EguA6CAjz4zazy+0RJ+FyfRdPe4TNEi4vnK8XsKNgWBIAZrpx7RKAIQsF8zzt80CuReBk4jRLyut58/JbuGbeFWMLAtrH7D6N2X4Uq+0wVtsRrLYjmF0NufX5wLeieLXLcGtXBMGgJggIOhy/0J+0pM63+0BaCcRYaa1pS7oDQkGWo+0Z2pN9fe0hU7GoKsSS6jCLc/UFS6rDLKoKE42Gwci1EpihCw7ZIxX5DRwNsj83GiS/5/x4MBpk7dxgJMji6lDfyb3oJJ8/wZdKMutxpsfhbI/D6W6Hsz1ZznQ7waXHoS3Zv27BVDCn3GZ+LhDMrwgxL24xu1xRHfOpimnKQxHidnnJjrGYNizilbVEo2NrKb0oAUApdRPwFcAE/kNr/fkBj78c+DRB95oL/K3W+uGRXlMCQH+e79Gd7cbxx3FVLNcJgsAQK2952ueBs4/y/cO/4GSikSXl9bwpFwRKMrGQk8LqOF4UCoJrI9Pddwzlc4sCQRAOvKqFpV0Z7wIVhwJLmYXroZqj860Enq+HHc0wJbhp7KansFr24lWvIFt35aBmTDF+utMexzoynMi1FBzvyHKiI0NDV9+cDABzy+0gENSEWVwdZnFtlKWz4tTEYygrPPoJbAYU+WVdn0NtmULtxv6mFEfbM4X3nlVmsa7oZL92TpSa2OT5N1ss4/qc7ckHgmwhGOTvt/S6/f6NKqCmzKTMNjFV0EphFrU8FLdCWLnFnIr3sVT//c2iAGQZweiTaCzGn1+ziqrY8LMNDmfcA4BSygQOAi8EGoAdwOu11vuK9ikHElprrZTaBPxUa71mpNeVADCY1ppep5fUeM+k5aRzhYKDWxw87fPg2cf478M/52SikcXldbx5+atKFwSKaY2RbC20EuRDgdl5vDCNqDZs3JolhVDg1S7DK5+HH6vNrS438f17wTCyvtEH+daCoUKB6/uF4jJdXGjmB+k52KYndh0FN4Xd9DT2mZ3Yp3diN+9BFQVT346RXfg8skuvJ7PoeTAOIz/EuWVdn1NdWU50ZINgkAsIJzqyJIumnC4LGSypDroQFtfGWFxbzuJZFdTPimMXrx2gNV6qi6NNHbmTfZr9TSkOt2WChX2AqogZDOnMzfOwbk6U2eUXuR6hxLK+Q9pLk/WyOJ6mNeHR1OPR3BtctydB+1auJiGoUfD8vlY+Txffp18tg6eLthfVPBTXMQA8+IHrWFR7/v+OLkYAuAr4hNb6xtz9DwNorT83wv7f0VqvHel1JQAML+tlSTrJ8SkQzNM6qA/IJoccxpsPAt8/8nNO9DayqKyON6+4hWvnXTnuCxDhOZidJ/t1IVjthzETrf1/BCuMH63Fj9XgxWrwY7PwYzX4sdrcJXc7WgPmxf9GYhpGodhwpFAwFE1fOMhXnefDQT4s+KUKC/kT/ukdhM7sxGreg/JdtDJxZ63FWXAZzvxtOHPWY7XsI3zsXsIn7sNItaPNENn655BZej3ZRdeiI5UXeDDiQmmtaUm4hVBwoiNbaEHID8WEoPm7rirMkpoY1WU2R5p7ONiSLixzXR4ycvM39J3w58XtSV18OVpaa9J+hrSXxjvHmgVhMzx+XQBaE62opbw8PqbP9WIEgFcDN2mt35m7/2bgCq31ewfs90rgc8Ac4MVa60dGel0JAOfm+i4pN0XaTZd2KuFivheMFnDSQz9cCAK/4HhvA4vKFvDm5bdw7fyrxj8IDKDSXVjtRzF6WzBS7RiJVsxUO0ayDSPRFmxLdw39c0SqCoHAKw4HhdtBeNCh8nFtVTCUgamM3GyGRt99wyzMoDgWvtaFYWnn5KRQZ3ZjNO5ANTyOOvt04YSv56zDr9+Grr8cf/4lEO77w6d1MHzO9TSO62KefZLQsXsJH7sXM3EWrUycBduCMLDkOvwSTTolSieR9QotBic6MhxvD263JV2W1YRZOzfK+tzJvr4yNLlP9r6L1X4Iq3kP2orizLsEP1434r9fz/dIeWkyfmbY+UQGGs8AAFBeWUu0bGz1TxcjAPwZcOOAAHC51vpvhtn/GuBjWusXDPHYu4B3ASxatGjriRMnxnrcM4qvfdJumqSbHL9CwSEmEhp4DA81Pc5/H/55IQi8afktbJ+AIDAizwmCQDIXDAqXdsyi20aybcjpS7VhgmH3jcnOT96Sm6ylb6y2iVZGYZx28fbi/fptN0y0GcaPVuNHq3LXfRcdqcAwQ6UPCE4KTj8JDY/Dqcfh7B7wneC45q6H+sth4eWw4NJ+J/xzKdQ4eD7+2Wcwj/ye0LF7sLpOoFG4czeTWXo9mSXPx6+oH9uxC5GjMt3YTU9jNe3Gbnoq6Joa0F3qRWfhztuCM+8SnLlbcGetBsMm4wXf9p0xTF400wPAeXUB5PY5BmzTWrcOt4+0AIxN2k2TclPjVyw4Qn0A5IPADr5/+Occ6z3FwrIFvGn5K7lu/nMmVxA4F61R2d7BQSHVCb6L0l4w5tn3gnoE7QW3c9fBbb9v+1D7aH/wdieNke4Kbg88JBQ6UjkgHPTd1tFqdKwWFa1BxWZhRioxDHNwQHCScHr3MCf8DcHJvv5yqLsEQqX9w+Z6Hl7rYTj0O8yj92K1HggOqXZNUDOw9Hq86uUlfU8xDWmN2XUCq+mp4GTftBur40jwkDJwa1bhzt2MM28LzpxNKDeFffbJ4NK0G7PnNAC+FSE1ay2p2etJztlIavZa/PP8nZ/pAcAiKAK8HmgkKAJ8g9b6maJ9VgBHckWAlwK/Aur1CC8uAeDCOL5Dyk2RcTOl7x44R30ABEHg4aYd/HdREHjj8lfw/HnPKenCQ9NSbtlUI9XR/5LsQKU6MNKd/bdneoZ+GcPqHxAiVZg9Z7Ga9xWa9L05a/HrtgbN+gsuwQjHMZRCYeTmsB/f0KY7T+If+j0c/j3mmd0AuFVLySy5nuzS63FnrS1Zl4simA5aDTFnf377iMc64j+jkf+N5aeQzk9UI9NJnyc3jd2yr+/bfdNTGOkOAPxQHGfu5uCEP3czzpyNIxaeOr6D092A1fQE0aa9RJv3EOk4gtI+GkWmehnJORtIzdlAcvYG3PK5Ix7ajA4AuTd5EfBlgmGA39Faf1Yp9W4ArfVtSqkPAm8BHCAFfECGAV4cvvZJuSlSbqr03QPnqA/Iv//DTTv4/pFfcLTnJPWx+Txn7lYiRpiQGSJshAibuYsRImSGiBTdHvi4bUyPIqOS8px+oUCligNC/7DgR6vJ1l2Ks+ASnHkbg5ES55A/QRq5UFAcDhTBdfF2hRrbf6PeJjh8L/rw7+DUDpT28Mvn57oJrseZu7nfULXiE3q/xXqKTuiGCoZVGar/KpMTzdd9VeKu7+dmt/MndnTHJGIkmoNv92d3B9/uWw+gdNA871Yuzp3st+DM3YxXveyccxrki/oyXgZ3iGZ+5aSItu4n1ryHaPNeoi37MHPdB05sdl8gmLORTNXSfr+HMz4AjAcJAKWltSbjZUi6ySH/AVyQc9QHQBAE/ti0kx8d/V+O9zaQHWMXhUIRNkOEDDt3nQ8HYcKmTcSMUF82j1UVy1hZsYS6svlTq+thGlEKFEbuWuXCQn7R5nyo6AsL+TChCAo6jaMPYBy5B078CeVl0ZFKdNUSqKxDVdajKuqgog4q6yG+AKzzHyc92fQFAz83FS540z0Y+B5W+0Gss33f7s3eoIlem2GcORuKvt1vgmhN8LyiUJfPm33Lh3t4voenPTztkvEdPN8f/eRbvke44yixlr1Em/cSa96LnWwBwLNjQbfBnI0k52zAn7uZ8tic0n0eA0gAECXleA5JN0lmFNN3nt8Lj1wfUMzXPlnfIeNlyfpZ0l6WrJcl42fJeFkyXoaM7xRtywTbc9vSfiZ4LLct2D9Lyk1zKnG6EDAiZpgVFUtYWbGUlRVLWFWxjEVlC6QbYgpR2QShk48SatiB1X0ao+csRu/ZwlwQENRHqLLZQRioqIPKOqioz13XQXxeMPvdFBWMF/fxc5NGeRq8qTpxFIDWhDqeJXLoTkKHf4NKBCdXXTYHveAS9PxLghqU2WtQZqjvBD+gFUdrjauDE73ru7i+h6vdEbtYPK3xdPBZen4wbNbTGt8fubPU6m0i1ryXaMteYs17CHccQ6GD0S11V5Be8SKyS64LRgqVkAQAMS483wuGEXrp0nUPaB0UmWUTE7YMuOd7nEg0cqj7GIe6j3Oo6xiHe46TzgWekGGzPL6YlZVLC8FgSflC7Ek0o6A4B9/DSLRg9pzF7DmD0X0Gq/dMcL/7DKq3uV8hpVYGqnxuEAYKAaG4BWFe34ySWoPvgpsBLwNuOnc7G4RcLxPcLzyev6SDfdyi5xTvE58Pi64MCizPYzTFSDytcbzgZFaYRGYSBgNDBbPbhRKnsQ/dhfnsr1HtR4I1CpZeC6tuhLqtwWc0TPeRr/3gJK+9wone015JW0fyw2Xz82t4uVUhh2qFMbK9RFv2U9G8h/jxP2D2nEabYTKLryWz4mayC59bktApAUCMK601aS9N0kniDVF9PiajqA+4mDzt05g4w8HuY4VgcLj7GIlcP5+tLJbGF+YCwVJWVi5lWflCQlP4W+OM5rm5gHAmuHSfwew9g9kdBAbV21xYewJAKxMVLgc3i/YywQiOMdKoYCpkKwxmGG2FwbRRXY0oLxO817yNsOiqIBDMv6Tk3RfFLQb5meYuVo2BAizTwDKCa9vpxjx0N+z/FTTuCnaquwzWvgRW3gTRqkGvUTjB564d3xu/Ic6jpMnPzhe0Fni5SbdsI0S5WYbZ9BThw3cRPnI3RroDPxQnu/QFpFfcjDP/stFPtTyABABx0WS9bDB6oFTdA6OoD5govvY5k2zmUPexfsGgx+kFgul7F5fXsSofCiqWsiy+iKgVmeAjFxfMczASzUEwyIUElenJnbBDaDMcrKJnhXNrtufuF9/O72eF0WY4dz+UmxNiiG+xbgb77F5CjTsJNe7Eaj4QVJ5bYfwFl+IvvAIWXYk5Zz3GOC3nW1x8GNQZXNioBAUYhsI2FZZhFK5xM3D0vuCkf+zBYJhpzXJY+1JY8xKorEfnm+K1l2vGz3/Dd6dUzUPYDFEZKjpB+y6cfAQO/BoO/T5oES2bDatfhF7zkmBuDaX6/YzFP27+nJvfZkarMMLnLtYdigQAMSYln2XwPOoDJpLWmqZUay4M9AWDzmywKJGBYlF5HSsrlrK6chmrKpaxvGIxkVIuBypmBJXpxT79ZBAIGnZidRwHwA/Hceq24izchl9/Oap6KaZh5iZ96psAqpQ0xcEgmM1xqH7xYAEbhW0pzNwJv9Avr/1gXon9v4JDd0O2F102G3/VzWRX34w3ayUeGl/7uRP/5P5bMFqDAkAxJw3H7of9v4bjDwRfhKoWByFo7Uugeum53yBSCaGxrachAUBckJIOI8zXBzipSR8Eimmtac10BIGg62ghFLRnOoEgFCwury+EgpUVSyUUiPNmJFqxG3cRatxFqGEnZm8TAF7ZHLL1l5Gtvwynbit+2azcEEcDS5kYuVCQH1mRPyWrohEXFzI80/V9PK2xDSM3iiNQ+AbfvB914FeYB3+L0duEtmNklm0nteoGnAWXjrnpe6oYMQAUS3cFLQLP3gknHwV00Bqw5iWw6kUQH2a+AQkAYqLl6wRSbqo0wwh9L0jDXja4PseCG5NNXyg4ysGuYxzsPsqzXUf7WgqUweKyOlblWglWVS5leXwxYakpEKOhNWZXA3bjTkINQSjIL43tVi8hW7eVbP02nAVb0OHz6x9WheFzqt8QTVW4b+SCBP1CQ7653tM+dJ8hdOh3RA7+Dqv9CNowyS68kvSqG8gsfi7YM6ebbNQBoFhvExz8bdAy0LQHUMEMnGteAitvCE76eRIAxGQyLqsR+n4QBnwnFwxKPFfBRRCEgvZCIMhfF4eCJeX1hUCwqmIZy+OLpNBQnJv2sVoPEWrYGbQSnHkK5aaDaW9nryFbdyl+2axcTUIkKDy0+m5rM4y2+4oStRUOKvDPo0VAZXoIH7mPyKHfYZ/ejULjzN1AetWNpJc/Hz1EMd9MMKYAUKzjOBy4Ew78Krht2LD0miAMLNsejFCRACAmG9d3g/kExmO6Yd/PtQxkp2wggNzSq+n2XCAIug8Odh2lywmm7zWVGYSCymWsqghaCRbE5lIVqpAZD8XwvCx20zOFQGA37Rty/YiRaGXkihf7QoHOhYbioKCtMEamh9DJx1C+g1u5kPSqG8msfCFepSzidMEBIE9raH4mKB48cBckmoMpjVe/CF7wCahaeN4vKQFAjLtxnW44T2tws7lA4E7KUQWjpbWmOd3WLxAc7D5Kd270AUDMjLIgNof5sbksyF3qctezIrUyu6Hoz8uiskmUl0E5aZSXATeLctMoN1PYjpcJ7hcuwXwFaqjtRdtQJpklV5NedSPu7DXjuiT2VFOyAFDM96BhR1AvcOQ+eO8OiNWc98tIABAXjda6EARKNp/A8G/W10LgOkHXwRQaOjRQfvTBsd5TnE42cSbZxOlUE6eTzZxJNuEWfZ62spgbnZ0LBnNYEJtXCAnzo7OlS0GIi2hcAkCxUDlEJmYeAJkaTYyaUoqYHSNmx8h4GVJOqrR1Av3fLDd5SghC5AKBC34m6D7QAy/jcxilopRiXmw282KzBz3maZ/WdBunk025S3NwnWpib8ezJL2+Nc8VilmRahZE57KgbB4LonMK4WBBbC7l9tjGEwshJsgEjqKQACDGJGyGCZthHN8JCga9bOnrBIopBZYNDDOBypChwAuCwyQPC6YymBudzdzobC6p3dDvMa01XU5PUThoyrUgNPNYy5OFoYp5ETNMVaiCylAFlXacylBwqQpVUGEH11WhCipCcapCccqs2LgvCSzEVJX1shzoOsy6qpXMipx/E/1kJwFAXBDbsKkMV47PugPnwzCAUZ7IBoUFL3c9+cKCUqpw0l5XtXLQ4yk3zelUEAhOJ5toy3TQle2hM9tNV7abE70NdDk9hXUSBjKUQWUuGORDQRAcKnLhoYKqcCWV4Soqw5VUWGWEDHNKzfMgxGgl3CTPdBxkT8ezPN2xn2c7j+Bol49u/hteu+ylE314JScBQJSEaZiUh8op02W5qT41OjcrWL7OxNd+YVvx/fy++etxN6awUBQOKG5Z0EFBT/72RRa1IiyPL2Z5fPGI+6W9DF3Zntylm06nm+5sD51OD11ObyE0HOttpDPbTU+2Z9j/FiEjRDwUJx4qJ27HidtlxO0yyu0yKqwyyq0YFVaMcruMuF3e97hVJqs0ikmlI9PFno4DucuzHOk+jo/GVCYrK5bwysU3cemsDTx37raJPtRxIQFAlJRSCluNfZ7zfsEBDRp8ghCRDw6udnNLg16EoYPnExZg6K4I/L7tkAsKuigwDLhddHVBFKBMMEwidoRIrJq5BPdRRu566GpvT3v0ZHvozHTSnemmM9NJV6aLnmxPcHF6CrfPplo43HWM7mw3aW/kxaHKrFhfWLDLKbf6QkJFqJzacDWzwjXMilQzK1Ijsy2KkskX4uZP+E93HOBU4jQAYSPE2qoVvHH5K9lUs5a1lSsKa4KMexHgBJIAICaV/FSno+mXDtYK7wsD+ctFaUUYzvkGhpHkQ0E+MAzaxoDtOjixn+PkPhqmMqkKV1EVrjqv5zm+Q0+2h95s76CgMPDSm+2hNXE6t08vrh4c6OJ2WSEQ1IZrmB2pYVakhlnhICDMjtRQYcdlTgUxiNaaE4lG9rTvZ0/Hs+zpOEBzug2AcquMDdWruanuWjZWr2FV5bIZuVz4zPuJxbSRb22wjf4tDgMDgavdqbkwSf6kNoVObrZhUxOpoeY8C6a01iTdJG2pNlqTLbQkm2lNNtOaaqE11Uprqo0jPSfpyHQNCni2YTMrXE1tJGg9GCok1ISrZ+Qf+GnLMMG0g1n1cisrep7Doc5D7Gndw9Nte9nbtp/u3IRcteEqNlav5bXVa9hYs5ql5Qul+BUJAGIasgwLa8Af+/y64/lQ4PjO1AwF05RSijK7jDK7jEUVi4beSWtcJ017qoWWZBOtyVZaU820pttpS3fQkm7jYPdR/tS8k6w/eCKpMiuGbVjYhoWlLGzDDm4bFraysAyz8JjVb7/gfv524TUGPF5hlzMnUsuc6CwqpVWidBRghnDw6XRSdLkJOrPddGQ66Mx00pHu4HDnYZ5pfabQBVVXXsdz6q5m4+yNbJq1ifmxeaihCn/94utJNjzoIpAAIGYE0zAxMQkX9Sn72sfzPRzfKYQDz/cmtgtBDE8prFCUOaFFzKksCgl+7g+474D20K5DT6aLlnQrremOICBk2unO9uLofAAMWoYc3yu0Ejm+S8ZJ4vhObj+vaD+3337nEjJs5kRmMSdaG4SCwu1ZzInOYk6kdsYvHOVpnx6nl85MFx3Z7sLIlQ4nKE7tzBWmdma66Mx00ls0k2YxS1ksqljETUtvYuOsjWyctZHaaO0Qe57jG39xQa/vAbmwoKygFW4aBgQJAGLGMpSBYRrYZl8Xgta60ELgahfHc8Z/1kNxYQwjV3sR/HdUYagoq6HCX8TyfDDI/2FXKqiRgNx1rmYiv734ehjBMrm54Jj7XXF8h65MF83JZpp7z9KcOE1zoonmVCs7Wp+ifYiui0o7XggDheui2zXhqinVTO3nT+i5kSYduRN6Z+G6h85sV3BSzwajUPwhwraBQUW4gupwNVWRKlZUrwhuh6uoilQValPy98usstK0tigF5hCnRDMMoQpw08FS51522oQBCQBCFFFKYZv2sKEgf5HugynAyI14oLTftJVShW6CYnNjc1lVvapvg9bgZsBN42STtKbbaU630pxqozndVrjdmDzLk23P9JvxEcBSJrMiNcyJ1DI7Uku5XUbEDBcuUTNMxIz0bbPyj0Vyj0eImuExD730fI8up6ffCXzgCT0/pLQr2zPsCR2CoruqUAVV4QoWltexMbKBqkgtVZHq4FJ0co+H4phqkg0XVQrsaHDRetqEAQkAQpzDUKHA135fS4HUFIihKAV2BOwIdqic+ZE48915w06i1Oskc6GgNQgI+et0G890HiThpkh7GZwh6htGYitrUDgYeLENi26nN9fk3k2X002Pkxj6x0IRt8sKM04uKqujsrpvlsnKcCVVuYmjqsJVVEaqsIxwUKw31DfsqWYahYFp8F9DiIvPUAYhM9RvYZ58oaHjOVN79IEoPdMCMw7hODgZcFPB6pdFyu0Y5fYilsWHKYLM8XyPtJ8h7WZIe32XlJfuux7wWDr3WPB4cN3t9NKcbiXtZnC1S9wupzJUwfKKRcEMkHZFbibIOFW5JvjKcCUV4UpM0yLoPjH7ukwucOjplFQcBnw/CANuesqEAQkAQpTIUIWGxcMR860FUmQ4w9nh4OJ7QRBw0uc1tbJpmJQZMcqs2IUfi2GAYRWdvPMn9FxdxTnqIUQRw4BQLLgUhwF36Gm4JwMJAEKMo+GGJHo6qD73tV+4LcWGM4xhBkvBhsqD1gA3FZwsxiMf9jvRW8Ft05KT+3iZImFAAoAQF1m+pSA0xDCwfBDIdyfkw4J0JUxzVii45E8WTio3FO085QsfDTOYJKfft3sxIQaFgVzQmwRhQAKAEJOIZVhYWDCgCDo/7XFx60E+KEiXwjRSfLJwndzJIt2/VSC/xoM54Bu9nOgnP8OAUFlwyYeBCVwgSwKAEFPAcNMeQ9Cl4Gu/UHQ48FK8AqOYQiw7uOh47ttirtBuOlTSi74wMIHkN0mIKS7fpWAz8iqMnu/1W1mxcBmwTQLDJJMfTihEiUkAEGKGyAeF0SgEghECQ742Ib+EsxBiapEAIIQYRCmFqUYfGAqtB1rj4xdqE4brkhBCTDwJAEKIC2Yoo/+89SPkhuLWheLgUNyiUPyYEGJ8SAAQQlxUY21dGKqwUYodhRg7CQBCiEltUOvCCAbWLnjaG7KeIR8WpH5BzGQSAIQQ08b5ti4AhRBQ3P0wKCigB42UkPAgpjoJAEKIGU0phUKNupWhWHF4KL7k6xmKbwsx2UgAEEKIMTqf8FAcBvKjJDzt4ft926VFYeIZykChUEphYAw5+dZ0IQFACCEugn61DMP0UBSPhhjYgtCvqyIXFCZjN0T+5Jm/XRyOio974LZi53r8XO+df9/CiTx3O38s+f8Wg/afYVMpSwAQQohJQimFpc7/z3I+CORrGNDBibO4VmHQdf527mRbOGnnTojFJ8fi7fkTZ/C//vuN9wk0f6zFoUBrXTjJi/MjAUAIIaa4wolacV4FkFNNcctC38YJOphpQCKTEEIIMQNJABBCCCFmIAkAQgghxAwkAUAIIYSYgSQACCGEEDOQBAAhhBBiBhpVAFBK3aSUelYpdVgp9aEhHn+jUurp3OVPSqnNpT9UIYQQQpTKOQOAUsoEvg7cDKwDXq+UWjdgt2PAtVrrTcCngW+V+kCFEEIIUTqjaQG4HDistT6qtc4CPwZeXryD1vpPWuuO3N1HgfrSHqYQQgghSmk0AaAOOFV0vyG3bTjvAH5zIQclhBBCiPE1mqmAh5poccjVGZRS1xEEgOcO8/i7gHcBLFq0aJSHKIQQQohSG00AaAAWFt2vB04P3EkptQn4D+BmrXXbUC+ktf4WufoApVSLUurEeR+xOJdZQOtEH8QMJZ/9xJHPfuLIZz9xVl/Ik0cTAHYAK5VSS4FG4HXAG4p3UEotAn4BvFlrfXA0b6y1nn2exypGQSm1U2t92UQfx0wkn/3Ekc9+4shnP3GUUjsv5PnnDABaa1cp9V7gboJVrL+jtX5GKfXu3OO3AR8DaoFv5FZrcuUXQgghhJi8RrUcsNb6LuCuAdtuK7r9TuCdpT00IYQQQowXmQlw+pE5GCaOfPYTRz77iSOf/cS5oM9eaT1kQb8QQgghpjFpARBCCCFmIAkAU5RSaqFS6j6l1H6l1DNKqf+T216jlPq9UupQ7rp6oo91ulJKmUqpJ5VSv87dl8/+IlBKVSmlfqaUOpD7/b9KPvuLQyn1vtzfm71Kqf9RSkXksx8fSqnvKKWalVJ7i7YN+1krpT6cW6/nWaXUjaN5DwkAU5cL/L3Wei1wJfDXuTUaPgTcq7VeCdybuy/Gx/8B9hfdl8/+4vgK8Fut9RpgM8F/A/nsx5lSqg74/4DLtNYbCEaFvQ757MfLfwE3Ddg25Ged+9v/OmB97jnfyK3jMyIJAFOU1vqM1vqJ3O0egj+CdQTrNHwvt9v3gFdMyAFOc0qpeuDFBJNf5clnP86UUhXANcB/Amits1rrTuSzv1gsIKqUsoAYwaRw8tmPA631g0D7gM3DfdYvB36stc5orY8BhwnW8RmRBIBpQCm1BLgEeAyYq7U+A0FIAOZM4KFNZ18G/gHwi7bJZz/+lgEtwHdz3S//oZQqQz77cae1bgT+BTgJnAG6tNa/Qz77i2m4z/p81+wBJABMeUqpcuDnwN9qrbsn+nhmAqXUS4BmrfWuiT6WGcgCLgW+qbW+BEggTc4XRa6/+eXAUmABUKaUetPEHpXIGfWaPcUkAExhSimb4OT/Q631L3Kbm5RS83OPzweaJ+r4prGrgZcppY4TLI/9fKXUD5DP/mJoABq01o/l7v+MIBDIZz/+XgAc01q3aK0dgunfn4N89hfTcJ/1qNbsGUgCwBSlgjmX/xPYr7X+16KH7gDemrv9VuCXF/vYpjut9Ye11vVa6yUEhTd/0Fq/Cfnsx53W+ixwSimVXwTlemAf8tlfDCeBK5VSsdzfn+sJao/ks794hvus7wBep5QK59btWQk8fq4Xk4mApiil1HOBh4A99PVD/yNBHcBPgUUE/2D/TGs9sJBElIhSajvwfq31S5RStchnP+6UUlsIii9DwFHg7QRfZuSzH2dKqU8CryUYhfQkwRTw5chnX3JKqf8BthOsttgEfBz4X4b5rJVSHwH+nOC/zd9qrX9zzveQACCEEELMPNIFIIQQQsxAEgCEEEKIGUgCgBBCCDEDSQAQQgghZiAJAEIIIcQMJAFACCGEmIEkAAghhBAzkAQAIYQQYgb6/wHsvdpSnMDsJAAAAABJRU5ErkJggg==\n",
+      "text/plain": [
+       "<Figure size 504x288 with 1 Axes>"
+      ]
+     },
+     "metadata": {
+      "needs_background": "light"
+     },
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "y_min = min([global_results[k][2].min() for k in global_results])\n",
+    "y_max = max([global_results[k][2].max() for k in global_results])\n",
+    "\n",
+    "fig2 = plt.figure(constrained_layout=True, figsize=(7, 4))\n",
+    "axes = plt.axes()\n",
+    "axes.set_xlim(2, 100)\n",
+    "#axes.set_ylim(y_min, y_max)\n",
+    "for k in global_results:\n",
+    "    plot_with_confidence_intervals(plt, global_results[k][1], global_results[k][2], global_results[k][3], label=k)\n",
+    "    #plt.plot(global_results_train_steps[k][1], global_results_train_steps[k][0], label=k)\n",
+    "plt.legend(loc=\"upper right\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python [conda env:prior-fitting]",
+   "language": "python",
+   "name": "conda-env-prior-fitting-py"
+  },
+  "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.9.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

prior-fitting/notebooks/FewShotOmniglot.ipynb

@@ -0,0 +1,168 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "976fbfea",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "sys.path.insert(0,'..')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "4b164f6b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from torch import nn\n",
+    "\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "\n",
+    "from train import train\n",
+    "import priors\n",
+    "import encoders\n",
+    "import positional_encodings\n",
+    "import utils\n",
+    "import bar_distribution\n",
+    "\n",
+    "\n",
+    "from samlib.utils import chunker"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "29d423b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mykwargs = \\\n",
+    "{\n",
+    " 'bptt': 5*5+1,\n",
+    "'nlayers': 6,\n",
+    " 'dropout': 0.0, 'steps_per_epoch': 100,\n",
+    " 'batch_size': 100}\n",
+    "mnist_jobs_5shot_pi_prior_search = [\n",
+    "    pretrain_and_eval( {'num_features': 28 * 28, 'fuse_x_y': False, 'num_outputs': 5,\n",
+    "                                            'translations': False, 'jonas_style': True}, priors.stroke.DataLoader, Losses.ce, enc, emsize=emsize, nhead=nhead, warmup_epochs=warmup_epochs, nhid=nhid, y_encoder_generator=encoders.get_Canonical(5), lr=lr, epochs=epochs, single_eval_pos_gen=mykwargs['bptt']-1,\n",
+    "                  extra_prior_kwargs_dict={'num_features': 28*28, 'fuse_x_y': False, 'num_outputs':5, 'only_train_for_last_idx': True,\n",
+    "                                          'min_max_strokes': (1,max_strokes), 'min_max_len': (min_len, max_len), 'min_max_width': (min_width, max_width), 'max_offset': max_offset, 'max_target_offset': max_target_offset},\n",
+    "                  **mykwargs)\n",
+    "    for max_strokes, min_len, max_len, min_width, max_width, max_offset, max_target_offset in random_hypers\n",
+    "    for enc in [encoders.Linear] for emsize in [1024] for nhead in [4] for nhid in [emsize*2] for warmup_epochs in [5] for lr in [.00001] for epochs in [128,1024] for _ in range(1)]\n",
+    "\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "deb93d1e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "@torch.inference_mode()\n",
+    "def get_acc(finetuned_model, eval_pos, device='cpu', steps=100, train_mode=False, **mykwargs):\n",
+    "    finetuned_model.to(device)\n",
+    "    finetuned_model.eval()\n",
+    "\n",
+    "    t_dl = priors.omniglot.DataLoader(steps, batch_size=1000, seq_len=mykwargs['bptt'], train=train_mode,\n",
+    "                                      **mykwargs['extra_prior_kwargs_dict'])\n",
+    "\n",
+    "    ps = []\n",
+    "    ys = []\n",
+    "    for x, y in tqdm(t_dl):\n",
+    "        p = finetuned_model(tuple(e.to(device) for e in x), single_eval_pos=eval_pos)\n",
+    "        ps.append(p)\n",
+    "        ys.append(y)\n",
+    "\n",
+    "    ps = torch.cat(ps, 1)\n",
+    "    ys = torch.cat(ys, 1)\n",
+    "\n",
+    "    def acc(ps, ys):\n",
+    "        return (ps.argmax(-1) == ys.to(ps.device)).float().mean()\n",
+    "\n",
+    "    a = acc(ps[eval_pos], ys[eval_pos]).cpu()\n",
+    "    print(a.item())\n",
+    "    return a\n",
+    "\n",
+    "\n",
+    "def train_and_eval(*args, **kwargs):\n",
+    "    r = train(*args, **kwargs)\n",
+    "    model = r[-1]\n",
+    "    acc = get_acc(model, -1, device='cuda:0', **kwargs).cpu()\n",
+    "    model.to('cpu')\n",
+    "    return [acc]\n",
+    "\n",
+    "def pretrain_and_eval(extra_prior_kwargs_dict_eval,*args, **kwargs):\n",
+    "    r = train(*args, **kwargs)\n",
+    "    model = r[-1]\n",
+    "    kwargs['extra_prior_kwargs_dict'] = extra_prior_kwargs_dict_eval\n",
+    "    acc = get_acc(model, -1, device='cuda:0', **kwargs).cpu()\n",
+    "    model.to('cpu')\n",
+    "    return r, acc"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "706ecbb7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "emsize = 1024\n",
+    "# mnist_jobs_5shot_pi[20].result()[-1].state_dict()\n",
+    "mykwargs = \\\n",
+    "    {'bptt': 5 * 5 + 1,\n",
+    "     'nlayers': 6,\n",
+    "     'nhead': 4, 'emsize': emsize,\n",
+    "     'encoder_generator': encoders.Linear, 'nhid': emsize * 2}\n",
+    "results = train_and_eval(priors.omniglot.DataLoader, Losses.ce, y_encoder_generator=encoders.get_Canonical(5),\n",
+    "                   load_weights_from_this_state_dict=mnist_jobs_5shot_pi_prior_search[67][0][-1].state_dict(), epochs=32, lr=.00001, dropout=dropout,\n",
+    "                   single_eval_pos_gen=mykwargs['bptt'] - 1,\n",
+    "                   extra_prior_kwargs_dict={'num_features': 28 * 28, 'fuse_x_y': False, 'num_outputs': 5,\n",
+    "                                            'translations': True, 'jonas_style': True},\n",
+    "                   batch_size=100, steps_per_epoch=200, **mykwargs)\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "611554b2",
+   "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.9.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

prior-fitting/notebooks/SetupForGPFittingExperiments.ipynb

@@ -0,0 +1,270 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "111c502f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "sys.path.insert(0,'..')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "e6b59ce3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from torch import nn\n",
+    "\n",
+    "\n",
+    "from train import train\n",
+    "import priors\n",
+    "import encoders\n",
+    "import positional_encodings\n",
+    "import utils\n",
+    "import bar_distribution\n",
+    "import transformer\n",
+    "\n",
+    "from samlib.utils import chunker"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "acf7423d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "kwargs = \\\n",
+    "{\n",
+    " 'nlayers': 6, \n",
+    " 'dropout': 0.0, 'steps_per_epoch': 100, \n",
+    "}\n",
+    "    \n",
+    "    \n",
+    "def train_and_compare_fast_gp_mix(*args, **kwargs):\n",
+    "    hps = kwargs['extra_prior_kwargs_dict']['hyperparameters']\n",
+    "    num_features = kwargs['extra_prior_kwargs_dict']['num_features']\n",
+    "    baseline_res = priors.fast_gp_mix.evaluate(\n",
+    "        *args[0].get_batch_method(10000,kwargs['bptt'],num_features, hyperparameters=hps),\n",
+    "        hyperparameters=hps, \n",
+    "        use_mse=Losses.mse == args[2])\n",
+    "    print(baseline_res, 'with fast_gp_mix')\n",
+    "    \n",
+    "    res = train(*args, **kwargs)\n",
+    "    return res, baseline_res\n",
+    "\n",
+    "def train_and_compare_fast_gp(*args, num_evals=1000, **kwargs):\n",
+    "    hps = kwargs['extra_prior_kwargs_dict']['hyperparameters']\n",
+    "    num_features = kwargs['extra_prior_kwargs_dict']['num_features']\n",
+    "    baseline_res = priors.fast_gp.evaluate(\n",
+    "        *args[0].get_batch_method(num_evals,kwargs['bptt'],num_features, hyperparameters=hps, device='cpu'),\n",
+    "        hyperparameters=hps, \n",
+    "        use_mse=Losses.mse == args[2], device='cpu')\n",
+    "    print(baseline_res, 'with fast_gp')\n",
+    "    \n",
+    "    res = train(*args, **kwargs)\n",
+    "    return res, baseline_res\n",
+    "\n",
+    "def train_and_compare_gp(*args, num_evals=10000, **kwargs):\n",
+    "    num_features = kwargs['extra_prior_kwargs_dict']['num_features']\n",
+    "    baseline_res = priors.gp.evaluate(\n",
+    "        *args[0].get_batch_method(num_evals,kwargs['bptt'],num_features),\n",
+    "        use_mse=Losses.mse == args[2])\n",
+    "    print(baseline_res, 'with fast_gp')\n",
+    "    \n",
+    "    res = train(*args, **kwargs)\n",
+    "    return res, baseline_res\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "da083e24",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import gpytorch\n",
+    "hps = {'noise': 1e-4, 'outputscale': 1., 'lengthscale': .6, 'fast_computations': (False,False,False)}\n",
+    "\n",
+    "import numpy as np, scipy.stats as st\n",
+    "\n",
+    "def compute_mean_and_conf_interval(accuracies, confidence=.95):\n",
+    "    accuracies = np.array(accuracies)\n",
+    "    n = len(accuracies)\n",
+    "    m, se = np.mean(accuracies, -1), st.sem(accuracies, -1)\n",
+    "    h = se * st.t.ppf((1 + confidence) / 2., n-1)\n",
+    "    return m, h\n",
+    "\n",
+    "\n",
+    "def bl(hps,bptt, num_evals=100, num_features=1, step_size=1, evals_per_batch=None, speedups=(False,False,False,False)):\n",
+    "    if evals_per_batch is None:\n",
+    "        evals_per_batch = num_evals\n",
+    "    else:\n",
+    "        assert num_evals%evals_per_batch == 0\n",
+    "    results = []\n",
+    "    for batch_i in range(num_evals//evals_per_batch):\n",
+    "        with gpytorch.settings.fast_computations(False,False,False):\n",
+    "            batch = priors.fast_gp.get_batch(evals_per_batch,bptt,num_features, hyperparameters=hps)\n",
+    "        with gpytorch.settings.fast_pred_var(speedups[0]), gpytorch.settings.fast_computations(*speedups[1:]):\n",
+    "            all_res, baseline_res,_ = priors.fast_gp.evaluate(\n",
+    "                    *batch,\n",
+    "                    hyperparameters=hps, step_size=step_size\n",
+    "                    )\n",
+    "        print(baseline_res, 'with fast_gp')\n",
+    "        \n",
+    "        results.append(all_res)\n",
+    "    all_results = torch.cat(results,1) # seq x batch_size\n",
+    "    return compute_mean_and_conf_interval(all_results) # mean array, var array\n",
+    "    \n",
+    "    \n",
+    "#settings = [{'num_evals':n,} for n in [100,1000]]\n",
+    "    \n",
+    "#js = [ex.submit(bl, hps, 2000, step_size=100, evals_per_batch=2, num_features=5, **kwargs) for kwargs in settings]\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8088aa12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# below you can simply replace the prior to priors.fast_gp_mix to do experiments over mixtures of GPs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "165e683c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "num_features = 5\n",
+    "hps = {'noise': 1e-4, 'outputscale': 1., 'lengthscale': .6, 'fast_computations': (False,False,False)}\n",
+    "ys = priors.fast_gp.get_batch(100000,20,num_features, hyperparameters=hps)[1]\n",
+    "fivefeature_jobs = [\n",
+    "    train(priors.fast_gp.DataLoader, bar_distribution.FullSupportBarDistribution(bar_distribution.get_bucket_limits(num_borders, ys=ys)), enc, emsize=emsize, nhead=nhead, warmup_epochs=warmup_epochs, y_encoder_generator=y_enc, pos_encoder_generator=pos_enc,\n",
+    "              batch_size=batch_size, scheduler=decay, extra_prior_kwargs_dict={'num_features': num_features, 'fuse_x_y': False, 'hyperparameters': hps},\n",
+    "              epochs=epochs, lr=lr, input_normalization=input_norm, bptt=2010, single_eval_pos_gen=single_eval_pos,aggregate_k_gradients=step_every, **kwargs) \n",
+    "    for enc in [encoders.Linear] for y_enc in [encoders.Linear] for emsize in [512] for nhead in [4] for nhid in [emsize*2] for epochs in [50*25,100*25,200*25,400*25] \n",
+    "    for warmup_epochs in [epochs//4] for input_norm in [False]\n",
+    "    for batch_size in [4] for step_every in [100//batch_size] for lr in [.0001,.0003,.001] for decay in [utils.get_cosine_schedule_with_warmup] for num_borders in [1000,10000] \n",
+    "    for single_eval_pos in [utils.get_weighted_single_eval_pos_sampler(2000)]\n",
+    "    for pos_enc in [positional_encodings.PositionalEncoding if single_eval_pos is None else positional_encodings.NoPositionalEncoding] \n",
+    "    for redo in range(1)\n",
+    "]\n",
+    "\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "15d01f3b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np, scipy.stats as st\n",
+    "\n",
+    "def compute_mean_and_conf_interval(accuracies, confidence=.95):\n",
+    "    accuracies = np.array(accuracies)\n",
+    "    n = len(accuracies)\n",
+    "    m, se = np.mean(accuracies), st.sem(accuracies)\n",
+    "    h = se * st.t.ppf((1 + confidence) / 2., n-1)\n",
+    "    return m, h\n",
+    "hps = {'noise': 1e-4, 'outputscale': 1., 'lengthscale': .6, 'fast_computations': (False,False,False)}\n",
+    "\n",
+    "@torch.inference_mode()\n",
+    "def run_test(model,device='cuda:0',step_size=100, start_pos=1, batch_size=1000, sub_batch_size=10, seq_len=2000):\n",
+    "    assert batch_size % sub_batch_size == 0\n",
+    "    model.to(device)\n",
+    "\n",
+    "    model.eval()\n",
+    "    nlls = []\n",
+    "    nll_confidences = []\n",
+    "    mses = []\n",
+    "    max_mses = []\n",
+    "    eval_positions = []\n",
+    "    \n",
+    "    def get_metrics(model, eval_pos, batch_size):\n",
+    "        x,y, target_y = priors.fast_gp.get_batch(batch_size=batch_size, seq_len=eval_pos+1, num_features=5,hyperparameters=hps, device=device)\n",
+    "        logits = model((x,y), single_eval_pos=eval_pos)\n",
+    "        if isinstance(model.criterion,nn.GaussianNLLLoss):\n",
+    "            nll = model.criterion(logits[0][...,0], target_y[eval_pos], var=logits[0][...,1].abs())\n",
+    "            return nll, 0., 0.\n",
+    "        means = model.criterion.mean(logits) # num_evals x batch_size\n",
+    "        maxs = (model.criterion.borders[logits.argmax(-1)] + model.criterion.borders[logits.argmax(-1)+1])/2\n",
+    "        mse = nn.MSELoss()\n",
+    "        nll = model.criterion(logits[0], target_y[eval_pos])\n",
+    "        return nll, mse(means[0], target_y[eval_pos]), mse(maxs[0], target_y[eval_pos])\n",
+    "        \n",
+    "    \n",
+    "    for eval_pos in range(start_pos, seq_len, step_size):\n",
+    "        eval_positions.append(eval_pos)\n",
+    "        print(eval_pos)\n",
+    "        \n",
+    "        nll = []\n",
+    "        mean_mse = []\n",
+    "        max_mse = []\n",
+    "        for i in range(batch_size//sub_batch_size):\n",
+    "            batch_nll, batch_mean_mse, batch_max_mse = get_metrics(model, eval_pos, sub_batch_size)\n",
+    "            nll.append(batch_nll)\n",
+    "            mean_mse.append(batch_mean_mse)\n",
+    "            max_mse.append(batch_max_mse)\n",
+    "        \n",
+    "        nll = torch.cat(nll)\n",
+    "        mean_mse = torch.tensor(mean_mse).mean()\n",
+    "        max_mse = torch.tensor(max_mse).mean()\n",
+    "        \n",
+    "        \n",
+    "        mses.append(mean_mse)\n",
+    "        max_mses.append(max_mse)\n",
+    "        nlls.append(nll.mean())\n",
+    "        nll_confidences.append(compute_mean_and_conf_interval(nll.to('cpu'))[1])\n",
+    "    return eval_positions, torch.stack(mses).to('cpu'), torch.stack(max_mses).to('cpu'), torch.stack(nlls).to('cpu'), torch.tensor(nll_confidences).to('cpu')\n",
+    "\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "755e88e4",
+   "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.9.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

prior-fitting/notebooks/Untitled.ipynb

@@ -0,0 +1,180 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "a873fcbb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "sys.path.insert(0,'..')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "56023c88",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random\n",
+    "\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "from torch import nn\n",
+    "from sklearn.gaussian_process import GaussianProcessRegressor\n",
+    "from sklearn.gaussian_process.kernels import RBF, DotProduct, WhiteKernel\n",
+    "from priors.utils import get_batch_to_dataloader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 68,
+   "id": "036c690b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_gp():\n",
+    "    gp = GaussianProcessRegressor(\n",
+    "            kernel=RBF(length_scale=.6, length_scale_bounds='fixed'),\n",
+    "            random_state=0, optimizer=None)\n",
+    "    return gp"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 77,
+   "id": "ff8a3cd1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "seq_len = 4\n",
+    "num_features = 10\n",
+    "x = torch.rand(seq_len, num_features)\n",
+    "gpr = get_gp()\n",
+    "y = gpr.sample_y(x, random_state=random.randint(0, 2 ** 32)).squeeze()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 78,
+   "id": "46fe34a9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[-0.29995838] [0.90399136]\n",
+      "[-0.1039504] [0.98874968]\n",
+      "[-0.03414801] [0.99876344]\n",
+      "[-0.01104748] [0.99986603]\n",
+      "[-0.00356252] [0.9999855]\n",
+      "[-0.00114827] [0.99999843]\n",
+      "[-0.00037014] [0.99999983]\n",
+      "[-0.00011934] [0.99999998]\n",
+      "[-3.8486538e-05] [1.]\n",
+      "[-1.24147253e-05] [1.]\n",
+      "[-4.00568455e-06] [1.]\n",
+      "[-1.2927993e-06] [1.]\n",
+      "[-4.17353027e-07] [1.]\n",
+      "[-1.34771328e-07] [1.]\n",
+      "[-4.35327732e-08] [1.]\n",
+      "[-1.40657691e-08] [1.]\n",
+      "[-4.54613576e-09] [1.]\n",
+      "[-1.46979425e-09] [1.]\n",
+      "[-4.75345491e-10] [1.]\n"
+     ]
+    }
+   ],
+   "source": [
+    "for num_copies in range(1,20):\n",
+    "    gp = get_gp()\n",
+    "    x_copied = x.tile((1,num_copies))\n",
+    "    gp.fit(x_copied[:-1],y[:-1])\n",
+    "    m,s = gp.predict(x_copied[-1].reshape(1,-1), return_std=True)\n",
+    "    print(m,s)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 79,
+   "id": "87752b3d",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([[1.        , 0.1047567 , 0.17720387, 0.33463634],\n",
+       "       [0.1047567 , 1.        , 0.14686013, 0.04858264],\n",
+       "       [0.17720387, 0.14686013, 1.        , 0.32035965],\n",
+       "       [0.33463634, 0.04858264, 0.32035965, 1.        ]])"
+      ]
+     },
+     "execution_count": 79,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "k = RBF(length_scale=.6, length_scale_bounds='fixed')\n",
+    "k(x)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 80,
+   "id": "6a409ae5",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([[1.00000000e+00, 2.41799081e-19, 5.26006251e-15, 9.26592960e-10],\n",
+       "       [2.41799081e-19, 1.00000000e+00, 1.48311381e-16, 1.10443925e-25],\n",
+       "       [5.26006251e-15, 1.48311381e-16, 1.00000000e+00, 4.04686299e-10],\n",
+       "       [9.26592960e-10, 1.10443925e-25, 4.04686299e-10, 1.00000000e+00]])"
+      ]
+     },
+     "execution_count": 80,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "k = RBF(length_scale=.6, length_scale_bounds='fixed')\n",
+    "k(x_copied)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "24141432",
+   "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.9.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

prior-fitting/positional_encodings.py

@@ -0,0 +1,70 @@
+import math
+
+import torch
+from torch import nn
+
+
+# Protocol for positonal encodings.
+# __init__(d_model, max_len=..[, more optionals])
+# forward(x: (seq_len, bs, d_model)) -> Tensor of shape (*x.shape[:2],d_model) containing pos. embeddings
+
+
+class NoPositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=None):
+        super(NoPositionalEncoding, self).__init__()
+        pass
+
+    def forward(self, x):
+        return x #* math.sqrt(x.shape[-1])
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = self.pe[:x.size(0), :] + x # * math.sqrt(x.shape[-1])
+        return x
+
+
+class LearnedPositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super(LearnedPositionalEncoding, self).__init__()
+        self.max_seq_len = max_len
+        #self.positional_embeddings = nn.Embedding(max_len, d_model)
+        self.positional_embeddings = nn.Parameter(torch.empty(max_len, d_model))
+        nn.init.normal_(self.positional_embeddings, mean=0, std=d_model ** -0.5)
+
+    def forward(self, x):
+        seq_len, bs, d_model = x.shape
+        assert seq_len <= len(self.positional_embeddings), 'seq_len can be at most max_len.'
+        pos_emb = self.positional_embeddings[:seq_len]
+        return pos_emb.unsqueeze(1).expand(seq_len, bs, d_model) + x #* math.sqrt(x.shape[-1])
+
+
+class PairedScrambledPositionalEncodings(LearnedPositionalEncoding):
+    # TODO check whether it is a problem to use the same perm. for full batch
+    def forward(self, x):
+        seq_len, bs, d_model = x.shape
+        assert seq_len <= len(self.positional_embeddings), 'seq_len can be at most max_len.'
+        assert len(self.positional_embeddings) % 2 == 0, 'Please specify an even max_len.'
+
+        paired_embs = self.positional_embeddings.view(len(self.positional_embeddings), -1, 2)
+        pos_emb = paired_embs[torch.randperm(len(paired_embs))].view(*self.positional_embeddings.shape)[:seq_len]
+
+        return pos_emb.unsqueeze(1).expand(seq_len, bs, d_model) + x #* math.sqrt(x.shape[-1])
+
+
+
+
+
+
+
+

prior-fitting/presentation/heatmap_bardistribution.py

@@ -0,0 +1,97 @@
+"""
+An example of how to use this:
+x ,y , y_target = priors.fast_gp.get_batch(1,100,num_features, hyperparameters=(1e-4,1.,.6), equidistant_x=True)
+fig, ax = pyplot.subplots(figsize=[10,10])
+plot_model_and_orig_curve(ax, SOME_MODEL, x, y, given_indices[10,40,60])
+
+Don't worry it is normal to be slow...
+"""
+import matplotlib.patches as patches
+import seaborn as sns
+import torch
+
+
+def add_rect(ax, coord, height, width, color):
+    rect = patches.Rectangle(coord, height, width, linewidth=1, edgecolor='none', facecolor=color)
+
+    # Add the patch to the Axes
+    ax.add_patch(rect)
+
+
+def heatmap_with_box_sizes(ax, data: torch.Tensor, x_starts, x_ends, y_starts, y_ends,
+                           palette=sns.color_palette("rocket", as_cmap=True), set_lims=True):
+    """
+    Beware all x and y arrays should be sorted from small to large and the data will appear in that same order: Small indexes map to lower x/y-axis values.
+    """
+    if set_lims:
+        ax.set_xlim(x_starts[0], x_ends[-1])
+        ax.set_ylim(y_starts[0], y_ends[-1])
+
+    data = (data - data.min()) / (data.max() - data.min())
+
+    for col_i, (col_start, col_end) in enumerate(zip(x_starts, x_ends)):
+        for row_i, (row_start, row_end) in enumerate(zip(y_starts, y_ends)):
+            add_rect(ax, (col_start, row_start), col_end - col_start, row_end - row_start,
+                     palette(data[row_i, col_i].item()))
+
+
+print(ax.get_ylim())
+
+
+def plot_bar_distribution(ax, x: torch.Tensor, bar_borders: torch.Tensor, predictions: torch.Tensor, **kwargs):
+    x = x.squeeze()
+    predictions = predictions.squeeze()
+    assert len(x.shape) == 1 and len(predictions.shape) == 2 and len(predictions) == len(x) and len(
+        bar_borders.shape) == 1 and len(bar_borders) - 1 == predictions.shape[1]
+
+    y_starts = bar_borders[:-1]
+    y_ends = bar_borders[1:]
+
+    x, order = x.sort(0)
+    print(x.shape, predictions.shape, order.shape)
+
+    predictions = predictions[order] / (bar_borders[1:] - bar_borders[:-1])
+    print(predictions.shape)
+
+    # assume x is sorted
+    x_starts = torch.cat([x[0].unsqueeze(0), (x[1:] + x[:-1]) / 2])
+    x_ends = torch.cat([(x[1:] + x[:-1]) / 2, x[-1].unsqueeze(0), ])
+
+    heatmap_with_box_sizes(ax, predictions.T, x_starts, x_ends, y_starts, y_ends, **kwargs)
+
+
+def plot_model_w_eval_pos(ax, model, x, y, single_eval_pos, softmax=False, min_max_y=None, **kwargs):
+    with torch.no_grad():
+        model.eval()
+        y_pred = model((x, y), single_eval_pos=single_eval_pos)
+        if softmax:
+            y_pred = y_pred.softmax(-1)
+    if min_max_y:
+        lowest_bar = torch.searchsorted(model.criterion.borders, min_max_y[0])
+        highest_bar = min(torch.searchsorted(model.criterion.borders, min_max_y[1]), len(model.criterion.borders))
+        borders = model.criterion.borders[lowest_bar:highest_bar]
+        y_pred = y_pred[..., lowest_bar:highest_bar - 1]
+    else:
+        borders = model.criterion.borders
+    plot_bar_distribution(ax, x[single_eval_pos:], borders, y_pred, **kwargs)
+
+
+def plot_model_and_orig_curve(ax, model, x, y, given_indices=[0]):
+    """
+    :param ax: A standard pyplot ax
+    :param model: A Transformer Model with `single_eval_pos`
+    :param x: A three-dimensional input tensor with x.shape[0]=1 and x.shape[2]=1
+    :param y: A two-dimensional tensor with y.shape[1]=0
+    :param given_indices: The indexes in y which should be given to the model (the training points)
+    :return:
+    """
+    x_winput = torch.cat([x[given_indices], x], 0)
+    y_winput = torch.cat([y[given_indices], y], 0)
+
+    ax.plot(x.squeeze(), y.squeeze(), color='grey')
+    ax.plot(x.squeeze()[given_indices], y.squeeze()[given_indices], 'o', color='black')
+    plot_model_w_eval_pos(ax, model, x_winput, y_winput, len(given_indices),
+                          min_max_y=(y.min() - .3, y.max() + .3), softmax=True,
+                          palette=sns.cubehelix_palette(start=2, rot=0, dark=0.4, light=1, as_cmap=True))
+
+

prior-fitting/priors/__init__.py

@@ -0,0 +1,4 @@
+from . import fast_gp, gp, ridge, stroke, fast_gp_mix, mlp, omniglot, binarized_regression, pyro
+
+
+

prior-fitting/priors/binarized_regression.py

@@ -0,0 +1,21 @@
+from . import fast_gp, fast_gp_mix
+from .utils import get_batch_to_dataloader
+
+def regression_prior_to_binary(get_batch_function):
+
+    def binarized_get_batch_function(*args, assert_on=False, **kwargs):
+        x, y, target_y = get_batch_function(*args, **kwargs)
+        if assert_on:
+            assert y is target_y, "y == target_y is assumed by this function"
+        y = y.sigmoid().bernoulli()
+        return x, y, y
+
+    return binarized_get_batch_function
+
+
+Binarized_fast_gp_dataloader = get_batch_to_dataloader(regression_prior_to_binary(fast_gp.get_batch))
+Binarized_fast_gp_dataloader.num_outputs = 1
+
+
+Binarized_fast_gp_mix_dataloader = get_batch_to_dataloader(regression_prior_to_binary(fast_gp_mix.get_batch))
+Binarized_fast_gp_mix_dataloader.num_outputs = 1

prior-fitting/priors/fast_gp.py

@@ -0,0 +1,130 @@
+import time
+
+import torch
+from torch import nn
+import gpytorch
+
+from .utils import get_batch_to_dataloader
+from utils import default_device
+from .utils import order_by_y, normalize_data, normalize_by_used_features_f, Binarize
+
+
+# We will use the simplest form of GP model, exact inference
+class ExactGPModel(gpytorch.models.ExactGP):
+    def __init__(self, train_x, train_y, likelihood):
+        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
+        self.mean_module = gpytorch.means.ConstantMean()
+        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
+
+    def forward(self, x):
+        mean_x = self.mean_module(x)
+        covar_x = self.covar_module(x)
+        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
+
+
+def get_model(x, y, hyperparameters):
+    likelihood = gpytorch.likelihoods.GaussianLikelihood(noise_constraint=gpytorch.constraints.GreaterThan(1.e-9))
+    model = ExactGPModel(x, y, likelihood)
+    model.likelihood.noise = torch.ones_like(model.likelihood.noise) * hyperparameters["noise"]
+    model.covar_module.outputscale = torch.ones_like(model.covar_module.outputscale) * hyperparameters["outputscale"]
+    model.covar_module.base_kernel.lengthscale = torch.ones_like(model.covar_module.base_kernel.lengthscale) * \
+                                                 hyperparameters["lengthscale"]
+    return model, likelihood
+
+
+@torch.no_grad()
+def get_batch(batch_size, seq_len, num_features, device=default_device, hyperparameters=None, equidistant_x=False, fix_x=None):
+    if isinstance(hyperparameters, (tuple, list)):
+        hyperparameters = {"noise": hyperparameters[0], "outputscale": hyperparameters[1], "lengthscale": hyperparameters[2]}
+    elif hyperparameters is None:
+        hyperparameters = {"noise": .1, "outputscale": .1, "lengthscale": .1}
+    with gpytorch.settings.fast_computations(*hyperparameters.get('fast_computations',(True,True,True))):
+        start = time.time()
+
+        assert not (equidistant_x and (fix_x is not None))
+        if equidistant_x:
+            assert num_features == 1
+            x = torch.linspace(0,1.,seq_len).unsqueeze(0).repeat(batch_size,1).unsqueeze(-1).to(device)
+        elif fix_x is not None:
+            assert fix_x.shape == (seq_len, num_features)
+            x = fix_x.unsqueeze(0).repeat(batch_size, 1, 1).to(device)
+        else:
+            x = torch.rand(batch_size, seq_len, num_features, device=device)
+        model, likelihood = get_model(x, torch.Tensor(), hyperparameters)
+        model.to(device)
+        # trained_model = ExactGPModel(train_x, train_y, likelihood).cuda()
+        # trained_model.eval()
+        with gpytorch.settings.prior_mode(True):
+            d = model(x)
+            d = likelihood(d)
+            sample = d.sample().transpose(0, 1)
+        #print(f'took {time.time() - start}')
+    return x.transpose(0, 1), sample, sample # x.shape = (T,B,H)
+
+# TODO: Reintegrate this code
+# num_features_used = num_features_used_sampler()
+# prior_outputscale = prior_outputscale_sampler()
+# prior_lengthscale = prior_lengthscale_sampler()
+#
+# x, sample = normalize_data(x), normalize_data(sample)
+#
+# if is_binary_classification:
+#     sample = (sample > torch.median(sample, dim=0)[0]).float()
+#
+# if normalize_by_used_features:
+#     x = normalize_by_used_features_f(x, num_features_used, num_features)
+#
+# # # if is_binary_classification and order_y:
+# # #     x, sample = order_by_y(x, sample)
+# #
+# # Append empty features if enabled
+# x = torch.cat([x, torch.zeros((x.shape[0], x.shape[1], num_features - num_features_used), device=device)], -1)
+
+DataLoader = get_batch_to_dataloader(get_batch)
+DataLoader.num_outputs = 1
+
+def get_model_on_device(x,y,hyperparameters,device):
+    model, likelihood = get_model(x, y, hyperparameters)
+    model.to(device)
+    return model, likelihood
+
+
+@torch.no_grad()
+def evaluate(x, y, y_non_noisy, use_mse=False, hyperparameters={}, get_model_on_device=get_model_on_device, device=default_device, step_size=1, start_pos=0):
+    start_time = time.time()
+    losses_after_t = [.0] if start_pos == 0 else []
+    all_losses_after_t = []
+
+    with gpytorch.settings.fast_computations(*hyperparameters.get('fast_computations',(True,True,True))), gpytorch.settings.fast_pred_var(False):
+        for t in range(max(start_pos, 1), len(x), step_size):
+            loss_sum = 0.
+            model, likelihood = get_model_on_device(x[:t].transpose(0, 1), y[:t].transpose(0, 1), hyperparameters, device)
+
+
+            model.eval()
+            # print([t.shape for t in model.train_inputs])
+            # print(x[:t].transpose(0,1).shape, x[t].unsqueeze(1).shape, y[:t].transpose(0,1).shape)
+            f = model(x[t].unsqueeze(1))
+            l = likelihood(f)
+            means = l.mean.squeeze()
+            varis = l.covariance_matrix.squeeze()
+            # print(l.variance.squeeze(), l.mean.squeeze(), y[t])
+
+            assert len(means.shape) == len(varis.shape) == 1
+            assert len(means) == len(varis) == x.shape[1]
+
+            if use_mse:
+                c = nn.MSELoss(reduction='none')
+                ls = c(means, y[t])
+            else:
+                ls = -l.log_prob(y[t].unsqueeze(1))
+
+            losses_after_t.append(ls.mean())
+            all_losses_after_t.append(ls.flatten())
+        return torch.stack(all_losses_after_t).to('cpu'), torch.tensor(losses_after_t).to('cpu'), time.time() - start_time
+
+if __name__ == '__main__':
+    hps = (.1,.1,.1)
+    for redo_idx in range(1):
+        print(
+            evaluate(*get_batch(1000, 10, hyperparameters=hps, num_features=10), use_mse=False, hyperparameters=hps))

prior-fitting/priors/fast_gp_mix.py

@@ -0,0 +1,307 @@
+import time
+import functools
+import random
+import math
+import traceback
+
+import torch
+from torch import nn
+import gpytorch
+from botorch.models import SingleTaskGP
+from botorch.models.gp_regression import MIN_INFERRED_NOISE_LEVEL
+from botorch.fit import fit_gpytorch_model
+from gpytorch.mlls import ExactMarginalLogLikelihood
+from gpytorch.likelihoods import GaussianLikelihood
+from gpytorch.priors.torch_priors import GammaPrior
+from gpytorch.constraints import GreaterThan
+
+
+from bar_distribution import BarDistribution
+from utils import default_device
+from .utils import get_batch_to_dataloader
+from . import fast_gp
+
+def get_model(x, y, hyperparameters: dict, sample=True):
+    aug_batch_shape = SingleTaskGP(x,y.unsqueeze(-1))._aug_batch_shape
+    noise_prior = GammaPrior(hyperparameters.get('noise_concentration',1.1), hyperparameters.get('noise_rate',0.05))
+    noise_prior_mode = (noise_prior.concentration - 1) / noise_prior.rate
+    likelihood = GaussianLikelihood(
+        noise_prior=noise_prior,
+        batch_shape=aug_batch_shape,
+        noise_constraint=GreaterThan(
+            MIN_INFERRED_NOISE_LEVEL,
+            transform=None,
+            initial_value=noise_prior_mode,
+        ),
+    )
+    model = SingleTaskGP(x, y.unsqueeze(-1),
+                         covar_module=gpytorch.kernels.ScaleKernel(
+                            gpytorch.kernels.MaternKernel(
+                                nu=hyperparameters.get('nu',2.5),
+                                ard_num_dims=x.shape[-1],
+                                batch_shape=aug_batch_shape,
+                                lengthscale_prior=gpytorch.priors.GammaPrior(hyperparameters.get('lengthscale_concentration',3.0), hyperparameters.get('lengthscale_rate',6.0)),
+                            ),
+                            batch_shape=aug_batch_shape,
+                            outputscale_prior=gpytorch.priors.GammaPrior(hyperparameters.get('outputscale_concentration',.5), hyperparameters.get('outputscale_rate',0.15)),
+                        ), likelihood=likelihood)
+
+    likelihood = model.likelihood
+    if sample:
+        sampled_model = model.pyro_sample_from_prior()
+        return sampled_model, sampled_model.likelihood
+    else:
+        assert not(hyperparameters.get('sigmoid', False)) and not(hyperparameters.get('y_minmax_norm', False)), "Sigmoid and y_minmax_norm can only be used to sample models..."
+        return model, likelihood
+
+
+@torch.no_grad()
+def get_batch(batch_size, seq_len, num_features, device=default_device, hyperparameters=None,
+              batch_size_per_gp_sample=None, num_outputs=1,
+              fix_to_range=None, equidistant_x=False):
+    '''
+    This function is very similar to the equivalent in .fast_gp. The only difference is that this function operates over
+    a mixture of GP priors.
+    :param batch_size:
+    :param seq_len:
+    :param num_features:
+    :param device:
+    :param hyperparameters:
+    :param for_regression:
+    :return:
+    '''
+    assert num_outputs == 1
+    hyperparameters = hyperparameters or {}
+    with gpytorch.settings.fast_computations(*hyperparameters.get('fast_computations',(True,True,True))):
+        batch_size_per_gp_sample = (batch_size_per_gp_sample or max(batch_size // 10,1))
+        assert batch_size % batch_size_per_gp_sample == 0
+
+        total_num_candidates = batch_size*(2**(fix_to_range is not None))
+        num_candidates = batch_size_per_gp_sample * (2**(fix_to_range is not None))
+        if equidistant_x:
+            assert num_features == 1
+            x = torch.linspace(0,1.,seq_len).unsqueeze(0).repeat(total_num_candidates,1).unsqueeze(-1)
+        else:
+            x = torch.rand(total_num_candidates, seq_len, num_features, device=device)
+        samples = []
+        for i in range(0,total_num_candidates,num_candidates):
+            num_of_dims ~ uniform
+            model, likelihood = get_model(x[i:i+num_candidates,...,:num_of_dims], torch.zeros(num_candidates,x.shape[1]), hyperparameters)
+            x[i:i + num_candidates, ..., num_of_dims:] = 0
+            x[i:i + num_candidates, ..., :num_of_dims] *= total_dims/num_of_dims
+            #print(model.covar_module.base_kernel.lengthscale)
+            model.to(device)
+            # trained_model = ExactGPModel(train_x, train_y, likelihood).cuda()
+            # trained_model.eval()
+            successful_sample = 0
+            throwaway_share = 0.
+            while successful_sample < 1:
+                with gpytorch.settings.prior_mode(True):
+                    d = model(x[i:i+num_candidates])
+                    d = likelihood(d)
+                    sample = d.sample() # bs_per_gp_s x T
+                    if hyperparameters.get('y_minmax_norm'):
+                        sample = ((sample - sample.min(1)[0]) / (sample.max(1)[0] - sample.min(1)[0]))
+                    if hyperparameters.get('sigmoid'):
+                        sample = sample.sigmoid()
+                    if fix_to_range is None:
+                        samples.append(sample.transpose(0, 1))
+                        successful_sample = True
+                        continue
+                    smaller_mask = sample < fix_to_range[0]
+                    larger_mask = sample >= fix_to_range[1]
+                    in_range_mask = ~ (smaller_mask | larger_mask).any(1)
+                    throwaway_share += (~in_range_mask[:batch_size_per_gp_sample]).sum()/batch_size_per_gp_sample
+                    if in_range_mask.sum() < batch_size_per_gp_sample:
+                        successful_sample -= 1
+                        if successful_sample < 100:
+                            print("Please change hyper-parameters (e.g. decrease outputscale_mean) it"
+                                  "seems like the range is set to tight for your hyper-parameters.")
+                        continue
+
+                    x[i:i+batch_size_per_gp_sample] = x[i:i+num_candidates][in_range_mask][:batch_size_per_gp_sample]
+                    sample = sample[in_range_mask][:batch_size_per_gp_sample]
+                    samples.append(sample.transpose(0, 1))
+                    successful_sample = True
+        if random.random() < .01:
+            print('throwaway share', throwaway_share/(batch_size//batch_size_per_gp_sample))
+
+        #print(f'took {time.time() - start}')
+        sample = torch.cat(samples, 1)
+        x = x.view(-1,batch_size,seq_len,num_features)[0]
+        # TODO think about enabling the line below
+        #sample = sample - sample[0, :].unsqueeze(0).expand(*sample.shape)
+        x = x.transpose(0,1)
+        assert x.shape[:2] == sample.shape[:2]
+        target_sample = sample
+    return x, sample, target_sample # x.shape = (T,B,H)
+
+
+class DataLoader(get_batch_to_dataloader(get_batch)):
+    num_outputs = 1
+    @torch.no_grad()
+    def validate(self, model, step_size=1, start_pos=0):
+        if isinstance(model.criterion, BarDistribution):
+            (x,y), target_y = self.gbm(**self.get_batch_kwargs, fuse_x_y=self.fuse_x_y)
+            model.eval()
+            losses = []
+            for eval_pos in range(start_pos, len(x), step_size):
+                logits = model((x,y), single_eval_pos=eval_pos)
+                means = model.criterion.mean(logits) # num_evals x batch_size
+                mse = nn.MSELoss()
+                losses.append(mse(means[0], target_y[eval_pos]))
+            model.train()
+            return torch.stack(losses)
+        else:
+            return 123.
+
+
+@torch.enable_grad()
+def get_fitted_model(x, y, hyperparameters, device):
+    # fit the gaussian process
+    model, likelihood = get_model(x,y,hyperparameters,sample=False)
+    #print(model.covar_module.base_kernel.lengthscale)
+    model.to(device)
+    mll = ExactMarginalLogLikelihood(likelihood, model)
+    model.train()
+    fit_gpytorch_model(mll)
+    #print(model.covar_module.base_kernel.lengthscale)
+    return model, likelihood
+
+
+evaluate = functools.partial(fast_gp.evaluate, get_model_on_device=get_fitted_model)
+
+def get_mcmc_model(x, y, hyperparameters, device, num_samples, warmup_steps):
+    from pyro.infer.mcmc import NUTS, MCMC
+    import pyro
+    x = x.to(device)
+    y = y.to(device)
+    model, likelihood = get_model(x, y, hyperparameters, sample=False)
+    model.to(device)
+
+
+    def pyro_model(x, y):
+        sampled_model = model.pyro_sample_from_prior()
+        _ = sampled_model.likelihood(sampled_model(x))
+        return y
+
+    nuts_kernel = NUTS(pyro_model, adapt_step_size=True)
+    mcmc_run = MCMC(nuts_kernel, num_samples=num_samples, warmup_steps=warmup_steps)
+    #print(x.shape)
+    mcmc_run.run(x, y)
+    model.pyro_load_from_samples(mcmc_run.get_samples())
+    model.eval()
+    # test_x = torch.linspace(0, 1, 101).unsqueeze(-1)
+    # test_y = torch.sin(test_x * (2 * math.pi))
+    # expanded_test_x = test_x.unsqueeze(0).repeat(num_samples, 1, 1)
+    # output = model(expanded_test_x)
+    #print(x.shape)
+    return model, likelihood
+    # output = model(x[-1].unsqueeze(1).repeat(1, num_samples 1))
+    # return output.mean
+
+
+
+
+def get_mean_logdensity(dists, x: torch.Tensor, full_range=None):
+    means = torch.cat([d.mean.squeeze() for d in dists], 0)
+    vars = torch.cat([d.variance.squeeze() for d in dists], 0)
+    assert len(means.shape) == 1 and len(vars.shape) == 1
+    dist = torch.distributions.Normal(means, vars.sqrt())
+    #logprobs = torch.cat([d.log_prob(x) for d in dists], 0)
+    logprobs = dist.log_prob(x)
+    if full_range is not None:
+        used_weight = 1. - (dist.cdf(torch.tensor(full_range[0])) + (1.-dist.cdf(torch.tensor(full_range[1]))))
+        if torch.isinf(-torch.log(used_weight)).any() or torch.isinf(torch.log(used_weight)).any():
+            print('factor is inf', -torch.log(used_weight))
+        logprobs -= torch.log(used_weight)
+    assert len(logprobs.shape) == 1
+    #print(logprobs)
+    return torch.logsumexp(logprobs, 0) - math.log(len(logprobs))
+
+
+def evaluate_(x, y, y_non_noisy, hyperparameters=None, device=default_device, num_samples=100, warmup_steps=300,
+              full_range=None, min_seq_len=0, use_likelihood=False):
+    with gpytorch.settings.fast_computations(*hyperparameters.get('fast_computations',(True,True,True))), gpytorch.settings.fast_pred_var(False):
+        x = x.to(device)
+        y = y.to(device)
+        start_time = time.time()
+        losses_after_t = [.0] if min_seq_len == 0 else []
+        all_losses = []
+
+        for t in range(max(min_seq_len,1), len(x)):
+            #print('Timestep', t)
+            loss_sum = 0.
+            step_losses = []
+            start_step = time.time()
+            for b_i in range(x.shape[1]):
+                done = 0
+                while done < 1:
+                    try:
+                        model, likelihood = get_mcmc_model(x[:t, b_i], y[:t, b_i], hyperparameters, device, num_samples=num_samples, warmup_steps=warmup_steps)
+                        model.eval()
+
+                        with torch.no_grad():
+                            dists = model(x[t, b_i, :].unsqueeze(
+                                0))  # TODO check what is going on here! Does the GP interpret the input wrong?
+                            if use_likelihood:
+                                dists = likelihood(dists)
+                            l = -get_mean_logdensity([dists], y[t, b_i], full_range)
+                        done = 1
+                    except Exception as e:
+                        done -= 1
+                        print('Trying again..')
+                        print(traceback.format_exc())
+                        print(e)
+                    finally:
+                        if done < -10:
+                            print('Too many retries...')
+                            exit()
+
+                step_losses.append(l.item())
+                #print('loss',l.item())
+                print(f'current average loss at step {t} is {sum(step_losses)/len(step_losses)} with {(time.time()-start_step)/len(step_losses)} s per eval.')
+                loss_sum += l
+
+            loss_sum /= x.shape[1]
+            all_losses.append(step_losses)
+            print(f'loss after step {t} is {loss_sum}')
+            losses_after_t.append(loss_sum)
+            print(f'losses so far {torch.tensor(losses_after_t)}')
+        return torch.tensor(losses_after_t), time.time() - start_time, all_losses
+
+
+
+
+
+if __name__ == '__main__':
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch_size', type=int)
+    parser.add_argument('--seq_len', type=int)
+    parser.add_argument('--min_seq_len', type=int, default=0)
+    parser.add_argument('--warmup_steps', type=int)
+    parser.add_argument('--num_samples', type=int)
+    parser.add_argument('--min_y', type=int)
+    parser.add_argument('--max_y', type=int)
+    parser.add_argument('--dim', type=int, default=1)
+    parser.add_argument('--use_likelihood', default=True, type=bool)
+    parser.add_argument('--device', default='cpu')
+    parser.add_argument('--outputscale_concentraion', default=2., type=float)
+    parser.add_argument('--noise_concentration', default=1.1, type=float)
+    parser.add_argument('--noise_rate', default=.05, type=float)
+
+    args = parser.parse_args()
+
+    print('min_y:', args.min_y)
+    full_range = (None if args.min_y is None else (args.min_y,args.max_y))
+
+    hps = {'outputscale_concentration': args.outputscale_concentraion, 'noise_concentration': args.noise_concentration,
+           'noise_rate': args.noise_rate, 'fast_computations': (False,False,False)}
+    x, y, _ = get_batch(args.batch_size, args.seq_len, args.dim, fix_to_range=full_range, hyperparameters=hps)
+    print('RESULT:', evaluate_(x, y, y, device=args.device, warmup_steps=args.warmup_steps,
+                               num_samples=args.num_samples, full_range=full_range, min_seq_len=args.min_seq_len,
+                               hyperparameters=hps, use_likelihood=args.use_likelihood))
+
+

prior-fitting/priors/gp.py

@@ -0,0 +1,70 @@
+import time
+import random
+
+import numpy as np
+import torch
+from torch import nn
+from sklearn.gaussian_process import GaussianProcessRegressor
+from sklearn.gaussian_process.kernels import RBF, DotProduct, WhiteKernel
+from .utils import get_batch_to_dataloader
+
+
+length_scale_sampling_gp = .6
+
+def get_gp(length_scale=None):
+    return GaussianProcessRegressor(
+        kernel=RBF(length_scale=length_scale or length_scale_sampling_gp, length_scale_bounds='fixed'),
+        random_state=0, optimizer=None)
+
+
+def get_batch(batch_size, seq_len, num_features, noisy_std=None):
+    # m = torch.normal(0.,.1,size=(batch_size,num_features))
+    # m2 = torch.rand(batch_size,num_features)
+    # b = 0 # torch.rand(batch_size)
+    x_t = torch.rand(batch_size, seq_len, num_features)
+    # gp_b = TensorGP(kernel=TensorRBF(noisy_std))
+    # y_t = gp_b.sample_from_GP_prior(x_t).detach()
+
+    gpr = get_gp(noisy_std)
+    y_t = torch.zeros(batch_size, seq_len)
+
+    for i in range(len(y_t)):
+        y_t[i] += gpr.sample_y(x_t[i], random_state=random.randint(0, 2 ** 32)).squeeze()
+    x, y = x_t.transpose(0, 1), y_t.transpose(0, 1)
+    # x, _ = torch.sort(x,dim=0)
+    return x, y, y
+
+
+DataLoader = get_batch_to_dataloader(get_batch)
+DataLoader.num_outputs = 1
+
+def evaluate(x, y, y_non_noisy, use_mse=False, length_scale=length_scale_sampling_gp):
+    start_time = time.time()
+    losses_after_t = [.0]
+    for t in range(1, len(x)):
+        loss_sum = 0.
+        for b_i in range(x.shape[1]):
+            gpr = get_gp(length_scale).fit(x[:t, b_i], y[:t, b_i])
+            means, stds = gpr.predict(x[t, b_i].unsqueeze(0), return_std=True)
+            assert len(means) == 1 == len(stds)
+            if use_mse:
+                c = nn.MSELoss()
+                l = c(torch.tensor(means), y[t, b_i].unsqueeze(-1))
+            else:
+                c = nn.GaussianNLLLoss(full=True)
+                l = c(torch.tensor(means), y[t, b_i].unsqueeze(-1),
+                      var=torch.tensor(stds) ** 2)
+            loss_sum += l
+
+
+        losses_after_t.append(loss_sum / x.shape[1])
+
+    return torch.tensor(losses_after_t), time.time()-start_time
+
+if __name__ == '__main__':
+    ls = .1
+    for alpha in set([ls, ls * 1.1, ls * .9]):
+        print(alpha)
+        for redo_idx in range(1):
+            print(
+                evaluate(*get_batch(1000, 10, noisy_std=ls, num_features=10), use_mse=False, length_scale=alpha))

prior-fitting/priors/mlp.py

@@ -0,0 +1,208 @@
+import random
+
+import torch
+from torch import nn
+
+import numpy as np
+
+from utils import default_device
+from .utils import get_batch_to_dataloader
+from .utils import order_by_y, normalize_data, normalize_by_used_features_f, Binarize
+from .utils import trunc_norm_sampler_f, beta_sampler_f, gamma_sampler_f, uniform_sampler_f, zipf_sampler_f, scaled_beta_sampler_f, uniform_int_sampler_f
+
+
+def canonical_pre_processing(x, canonical_args):
+    assert x.shape[2] == len(canonical_args)
+    ranges = [torch.arange(num_classes).float() if num_classes is not None else None for num_classes in canonical_args]
+    for feature_dim, rang in enumerate(ranges):
+        if rang is not None:
+            x[:, :, feature_dim] = (x[:, :, feature_dim] - rang.mean()) / rang.std()
+    return x
+
+
+DEFAULT_NUM_LAYERS = 2
+DEFAULT_HIDDEN_DIM = 100
+DEFAULT_ACTIVATION_MODULE = torch.nn.ReLU
+DEFAULT_INIT_STD = .1
+DEFAULT_HIDDEN_NOISE_STD = .1
+DEFAULT_FIXED_DROPOUT = 0.
+DEFAULT_IS_BINARY_CLASSIFICATION = False
+
+
+class GaussianNoise(nn.Module):
+    def __init__(self, std):
+        super().__init__()
+        self.std = std
+
+    def forward(self, x):
+        return x + torch.normal(torch.zeros_like(x), self.std)
+
+
+def causes_sampler_f(num_causes_sampler):
+    num_causes = num_causes_sampler()
+    means = np.random.normal(0, 1, (num_causes))
+    std = np.abs(np.random.normal(0, 1, (num_causes)) * means)
+    return means, std
+
+def categorical_features_sampler(max_features):
+    features = []
+    ordinal = []
+    num_categorical_features_sampler = scaled_beta_sampler_f(0.5, .8, max_features, 0)
+    is_ordinal_sampler = lambda : random.choice([True, False])
+    classes_per_feature_sampler = scaled_beta_sampler_f(0.1, 2.0, 10, 1)
+    classes_per_feature_sampler_ordinal = scaled_beta_sampler_f(0.1, 2.0, 200, 1)
+    for i in range(0, num_categorical_features_sampler()):
+        ordinal_s = is_ordinal_sampler()
+        ordinal.append(ordinal_s)
+        classes = classes_per_feature_sampler_ordinal() if ordinal_s else classes_per_feature_sampler()
+        features.append(np.random.rand(classes))
+    return features, ordinal
+
+
+def get_batch(batch_size, seq_len, num_features, device=default_device, hyperparameters=(DEFAULT_NUM_LAYERS, DEFAULT_HIDDEN_DIM, DEFAULT_ACTIVATION_MODULE, DEFAULT_INIT_STD, DEFAULT_HIDDEN_NOISE_STD, DEFAULT_FIXED_DROPOUT, DEFAULT_IS_BINARY_CLASSIFICATION),
+              batch_size_per_gp_sample=None, num_outputs=1, canonical_args=None, sampling='normal'):
+    assert num_outputs == 1
+    num_layers_sampler, hidden_dim_sampler, activation_module, init_std_sampler, noise_std_sampler, dropout_prob_sampler, is_binary_classification, num_features_used_sampler, causes_sampler, is_causal, pre_sample_causes, pre_sample_weights, y_is_effect, order_y, normalize_by_used_features, categorical_features_sampler, nan_prob = hyperparameters
+
+    # if is_binary_classification:
+    #     sample_batch_size = 100*batch_size
+    # else:
+    sample_batch_size = batch_size
+
+    # if canonical_args is not None:
+    #     assert len(canonical_args) == num_causes
+    #     # should be list of [None, 2, 4] meaning scalar parameter, 2 classes, 4 classes
+    #
+    #     for feature_idx, num_classes in enumerate(canonical_args):
+    #         if num_classes is not None:
+    #             causes[:,:,feature_idx] = torch.randint(num_classes, (seq_len, sample_batch_size))
+    #
+    #     causes = canonical_pre_processing(causes, canonical_args)
+
+    batch_size_per_gp_sample = batch_size_per_gp_sample or sample_batch_size // 8
+    assert sample_batch_size % batch_size_per_gp_sample == 0, 'Please choose a batch_size divisible by batch_size_per_gp_sample.'
+    num_models = sample_batch_size // batch_size_per_gp_sample
+    # standard kaiming uniform init currently...
+
+    def get_model():
+        class MLP(torch.nn.Module):
+            def __init__(self):
+                super(MLP, self).__init__()
+
+                self.dropout_prob = dropout_prob_sampler()
+                self.noise_std = noise_std_sampler()
+                self.init_std = init_std_sampler()
+                self.num_features_used = num_features_used_sampler()
+                self.categorical_features, self.categorical_features_is_ordinal = categorical_features_sampler(self.num_features_used)
+                if is_causal:
+                    self.causes = causes_sampler() if is_causal else self.num_features_used
+                    self.causes = (torch.tensor(self.causes[0], device=device).unsqueeze(0).unsqueeze(0).tile((seq_len,1,1)), torch.tensor(self.causes[1], device=device).unsqueeze(0).unsqueeze(0).tile((seq_len,1,1)))
+                    self.num_causes = self.causes[0].shape[2]
+                else:
+                    self.num_causes = self.num_features_used
+                self.num_layers = num_layers_sampler()
+                self.hidden_dim = hidden_dim_sampler()
+
+                if is_causal:
+                    self.hidden_dim = max(self.hidden_dim, 2 * self.num_features_used+1)
+
+                #print('cat', self.categorical_features, self.categorical_features_is_ordinal, self.num_features_used)
+
+                assert(self.num_layers > 2)
+
+                self.layers = [nn.Linear(self.num_causes, self.hidden_dim, device=device)]
+                self.layers += [module for layer_idx in range(self.num_layers-1) for module in [
+                        nn.Sequential(*[
+                            activation_module()
+                            , nn.Linear(self.hidden_dim, num_outputs if layer_idx == self.num_layers - 2 else self.hidden_dim, device=device)
+                            , GaussianNoise(torch.abs(torch.normal(torch.zeros((num_outputs if layer_idx == self.num_layers - 2 else self.hidden_dim),device=device), self.noise_std))) if pre_sample_weights else GaussianNoise(self.noise_std)
+                        ])
+                    ]]
+                self.layers = nn.Sequential(*self.layers)
+
+                self.binarizer = Binarize() if is_binary_classification else lambda x : x
+
+                # Initialize Model parameters
+                for i, p in enumerate(self.layers.parameters()):
+                    dropout_prob = self.dropout_prob if i > 0 else 0.0
+                    nn.init.normal_(p, std=self.init_std / (1. - dropout_prob))
+                    with torch.no_grad():
+                        p *= torch.bernoulli(torch.zeros_like(p) + 1. - dropout_prob)
+
+            def forward(self):
+                if sampling == 'normal':
+                    if is_causal and pre_sample_causes:
+                        causes = torch.normal(self.causes[0], self.causes[1].abs()).float()
+                    else:
+                        causes = torch.normal(0., 1., (seq_len, 1, self.num_causes), device=device).float()
+                elif sampling == 'uniform':
+                    causes = torch.rand((seq_len, 1, self.num_causes), device=device)
+                else:
+                    raise ValueError(f'Sampling is set to invalid setting: {sampling}.')
+
+                outputs = [causes]
+                for layer in self.layers:
+                    outputs.append(layer(outputs[-1]))
+                outputs = outputs[2:]
+
+                if is_causal:
+                    outputs_flat = torch.cat(outputs, -1)
+                    random_perm = torch.randperm(outputs_flat.shape[-1]-1, device=device)
+                    random_idx_y = [-1] if y_is_effect else random_perm[0:num_outputs]
+                    y = outputs_flat[:, :, random_idx_y]
+
+                    random_idx = random_perm[num_outputs:num_outputs + self.num_features_used]
+                    x = outputs_flat[:, :, random_idx]
+                else:
+                    y = outputs[-1][:, :, :]
+                    x = causes
+
+                if len(self.categorical_features) > 0:
+                    random_perm = torch.randperm(x.shape[-1], device=device)
+                    for i, (categorical_feature, is_ordinal) in enumerate(zip(self.categorical_features, self.categorical_features_is_ordinal)):
+                        idx = random_perm[i]
+                        temp = normalize_data(x[:, :, idx])
+                        if is_ordinal:
+                            x[:, :, idx] = (temp > (torch.tensor(categorical_feature, device=device, dtype=torch.float32).unsqueeze(-1).unsqueeze(-1) - 0.5)).sum(axis=0)
+                        else:
+                            x[:, :, idx] = (temp > (torch.tensor(categorical_feature, device=device,
+                                                                dtype=torch.float32).unsqueeze(-1).unsqueeze(-1) - 0.5)).sum(
+                                axis=0) * (127 * len(categorical_feature) + 1) % len(categorical_feature)
+
+
+                # if nan_prob > 0:
+                #     nan_value = random.choice([-999,-1,0, -10])
+                #     x[torch.rand(x.shape, device=device) > (1-nan_prob)] = nan_value
+
+                x, y = normalize_data(x), normalize_data(y)
+
+                # Binarize output if enabled
+                y = self.binarizer(y)
+
+                if normalize_by_used_features:
+                    x = normalize_by_used_features_f(x, self.num_features_used, num_features)
+
+                if is_binary_classification and order_y:
+                    x, y = order_by_y(x,y)
+
+                # Append empty features if enabled
+                x = torch.cat([x, torch.zeros((x.shape[0], x.shape[1], num_features - self.num_features_used), device=device)], -1)
+
+                return x, y
+
+        return MLP()
+
+    models = [get_model() for _ in range(num_models)]
+
+    sample = sum([[model() for _ in range(0,batch_size_per_gp_sample)] for model in models],[])
+
+    x, y = zip(*sample)
+    y = torch.cat(y, 1).squeeze(-1).detach()
+    x = torch.cat(x, 1).detach()
+
+    return x, y, y
+
+
+DataLoader = get_batch_to_dataloader(get_batch)
+DataLoader.num_outputs = 1
+

prior-fitting/priors/omniglot.py

@@ -0,0 +1,98 @@
+import math
+import random
+import torch
+from torch.utils import data
+from torchvision import transforms
+import numpy as np
+
+from datasets import omniglotNshot
+import utils
+
+
+def _compute_maxtranslations(single_image_tensor, dim, background):
+    assert len(single_image_tensor.shape) == 2
+    content_rows = ((single_image_tensor == background).all(dim=1 - dim) == False).nonzero()
+    begin, end = content_rows[0], content_rows[-1]
+    return torch.cat([-begin, single_image_tensor.shape[dim] - end - 1]).cpu().tolist()
+
+
+def compute_maxtranslations_x_y(single_image_tensor, background):
+    return _compute_maxtranslations(single_image_tensor, 1, background), _compute_maxtranslations(single_image_tensor,
+                                                                                                  0, background)
+
+
+def translate(img, trans_x, trans_y):
+    return transforms.functional.affine(img.unsqueeze(0), angle=0.0, translate=[trans_x, trans_y], scale=1.0,
+                                        interpolation=transforms.InterpolationMode.NEAREST, shear=[0.0, 0.0],
+                                        fill=0.).squeeze(0)
+
+def translate_omniglot(image_tensor, background=0.):
+    flat_image_tensor = image_tensor.view(-1, *image_tensor.shape[-2:])
+    for i, image in enumerate(flat_image_tensor):
+        max_x, max_y = compute_maxtranslations_x_y(image, background)
+        flat_image_tensor[i] = translate(image, random.randint(*max_x), random.randint(*max_y))
+    return flat_image_tensor.view(*image_tensor.shape)
+
+
+class DataLoader(data.DataLoader):
+    def __init__(self, num_steps, batch_size, seq_len, num_features, num_outputs, num_classes_used=1200, fuse_x_y=False, train=True, translations=True, jonas_style=False):
+        # TODO position before last is predictable by counting..
+        utils.set_locals_in_self(locals())
+        assert not fuse_x_y, 'So far don\' support fusing.'
+        imgsz = math.isqrt(num_features)
+        assert imgsz * imgsz == num_features
+        assert ((seq_len-1) // num_outputs) * num_outputs == seq_len - 1
+        if jonas_style:
+            self.d = omniglotNshot.OmniglotNShotJonas('omniglot', batchsz=batch_size, n_way=num_outputs,
+                                                 k_shot=((seq_len - 1) // num_outputs),
+                                                 k_query=1, imgsz=imgsz)
+        else:
+            self.d = omniglotNshot.OmniglotNShot('omniglot', batchsz=batch_size, n_way=num_outputs,
+                                                 k_shot=((seq_len - 1) // num_outputs),
+                                                 k_query=1, imgsz=imgsz, num_train_classes_used=num_classes_used)
+
+
+    def __len__(self):
+        return self.num_steps
+
+    def __iter__(self):
+        # Eval at pos
+        def t(x, y, x_q, y_q):
+            x = np.concatenate([x,x_q[:,:1]], 1)
+            y = np.concatenate([y,y_q[:,:1]], 1)
+            y = torch.from_numpy(y).transpose(0, 1)
+            target_y = y.clone().detach()
+            target_y[:-1] = -100
+            x = torch.from_numpy(x)
+            if self.translations and self.train:
+                x = translate_omniglot(x)
+            image_tensor = x.view(*x.shape[:2], -1).transpose(0, 1), y
+            return image_tensor, target_y
+
+        return (t(*self.d.next(mode='train' if self.train else 'test')) for _ in range(self.num_steps))
+
+    @torch.no_grad()
+    def validate(self, finetuned_model, eval_pos=-1):
+        finetuned_model.eval()
+        device = next(iter(finetuned_model.parameters())).device
+
+        if not hasattr(self, 't_dl'):
+            self.t_dl = DataLoader(num_steps=self.num_steps, batch_size=self.batch_size, seq_len=self.seq_len,
+                                   num_features=self.num_features, num_outputs=self.num_outputs, fuse_x_y=self.fuse_x_y,
+                                   train=False)
+
+        ps = []
+        ys = []
+        for x,y in self.t_dl:
+            p = finetuned_model(tuple(e.to(device) for e in x), single_eval_pos=eval_pos)
+            ps.append(p)
+            ys.append(y)
+
+        ps = torch.cat(ps,1)
+        ys = torch.cat(ys,1)
+
+        def acc(ps,ys):
+            return (ps.argmax(-1)==ys.to(ps.device)).float().mean()
+
+        a = acc(ps[eval_pos], ys[eval_pos]).cpu()
+        return a

prior-fitting/priors/prior.py

@@ -0,0 +1,12 @@
+from torch.utils.data import DataLoader
+
+
+class PriorDataLoader(DataLoader):
+    pass
+    # init accepts num_steps as first argument
+
+    # has two attributes set on class or object level:
+    # num_features: int and
+    # num_outputs: int
+    # fuse_x_y: bool
+    # Optional: validate function that accepts a transformer model

prior-fitting/priors/pyro.py

@@ -0,0 +1,39 @@
+import random
+
+import torch
+from torch import nn
+
+from utils import default_device
+from .utils import get_batch_to_dataloader
+
+
+def get_batch(batch_size, seq_len, batch_size_per_gp_sample=None, **config):
+    batch_size_per_gp_sample = batch_size_per_gp_sample or batch_size // 16
+    assert batch_size % batch_size_per_gp_sample == 0, 'Please choose a batch_size divisible by batch_size_per_gp_sample.'
+    num_models = batch_size // batch_size_per_gp_sample
+    # standard kaiming uniform init currently...
+
+    models = [config['model']() for _ in range(num_models)]
+
+    sample = sum([[model(seq_len=seq_len) for _ in range(0,batch_size_per_gp_sample)] for model in models],[])
+
+    def normalize_data(data):
+        mean = data.mean(0)
+        std = data.std(0) + .000001
+        eval_xs = (data - mean) / std
+
+        return eval_xs
+
+    x, y = zip(*sample)
+
+    y = torch.stack(y, 1).squeeze(-1).detach()
+    x = torch.stack(x, 1).detach()
+
+    x, y = normalize_data(x), y
+
+    return x, y, y
+
+
+DataLoader = get_batch_to_dataloader(get_batch)
+DataLoader.num_outputs = 1
+

prior-fitting/priors/ridge.py

@@ -0,0 +1,38 @@
+import random
+import time
+
+import numpy as np
+import torch
+from torch import nn
+from sklearn.linear_model import Ridge
+from .utils import get_batch_to_dataloader
+
+def get_batch(batch_size, seq_len, num_features, noisy_std = .1):
+    m = torch.normal(0., .1, size=(batch_size,num_features))
+    b = 0 # torch.rand(batch_size)
+    x = torch.rand(seq_len, batch_size,num_features)
+    y_non_noisy = torch.einsum('bf,tbf->tb',m,x)
+    y = y_non_noisy + torch.normal(torch.zeros_like(y_non_noisy),noisy_std) # noisy_std is alpha
+    return x, y, y_non_noisy
+
+DataLoader = get_batch_to_dataloader(get_batch)
+DataLoader.num_outputs = 1
+
+
+def evaluate(x,y,y_non_noisy, alpha=0.):
+    start_time = time.time()
+    losses_after_t = [.0]
+    for t in range(1,len(x)):
+        loss_sum = 0.
+        for b_i in range(x.shape[1]):
+            clf = Ridge(alpha=alpha)
+            clf.fit(x[:t,b_i],y[:t,b_i])
+            y_ = clf.predict(x[t,b_i].unsqueeze(0))
+            l = nn.MSELoss()(y_non_noisy[t,b_i].unsqueeze(0),torch.tensor(y_))
+            loss_sum += l
+        losses_after_t.append(loss_sum/x.shape[1])
+    return torch.tensor(losses_after_t), time.time()-start_time
+
+if __name__ == '__main__':
+    for alpha in [.001,.01,.5,1.]:
+        print(alpha, evaluate(*get_batch(1000,10,noisy_std=.01),alpha=alpha))

prior-fitting/priors/stroke.py

@@ -0,0 +1,143 @@
+from PIL import Image, ImageDraw, ImageFilter
+import random
+import math
+
+import torch
+import numpy as np
+from .utils import get_batch_to_dataloader
+
+def mnist_prior(num_classes=2, size=28, min_max_strokes=(1,3), min_max_len=(5/28,20/28), min_max_start=(2/28,25/28),
+                min_max_width=(1/28,4/28), max_offset=4/28, max_target_offset=2/28):
+    classes = []
+    for i in range(num_classes):
+        num_strokes = random.randint(*min_max_strokes)
+        len_strokes = [random.randint(int(size * min_max_len[0]), int(size * min_max_len[1])) for i in range(num_strokes)]
+        stroke_start_points = [
+            (random.randint(int(size * min_max_start[0]), int(size * min_max_start[1])), random.randint(int(size * min_max_start[0]), int(size * min_max_start[1]))) for i in
+            range(num_strokes)]
+        stroke_directions = []
+        # i = Image.fromarray(np.zeros((28,28),dtype=np.uint8))
+        # draw = ImageDraw.Draw(i)
+        for i in range(num_strokes):
+            sp, length = stroke_start_points[i], len_strokes[i]
+            counter = 0
+            while True:
+                if counter % 3 == 0:
+                    length = random.randint(int(size * min_max_len[0]), int(size * min_max_len[1]))
+                    sp = (
+                    random.randint(int(size * min_max_start[0]), int(size * min_max_start[1])), random.randint(int(size * min_max_start[0]), int(size * min_max_start[1])))
+                    stroke_start_points[i], len_strokes[i] = sp, length
+                radians = random.random() * 2 * math.pi
+                x_vel = math.cos(radians) * length
+                y_vel = math.sin(radians) * length
+                new_p = (sp[0] + x_vel, sp[1] + y_vel)
+                # print(math.degrees(radians),sp,new_p)
+                if not any(n > size - 1 or n < 0 for n in new_p):
+                    break
+                counter += 1
+            stroke_directions.append(radians)
+            # print([round(x) for x in sp+new_p])
+            # draw.line([round(x) for x in sp+new_p], fill=128, width=3)
+        classes.append((len_strokes, stroke_start_points, stroke_directions))
+
+    generator_functions = []
+    for c in classes:
+        def g(c=c):
+            len_strokes, stroke_start_points, stroke_directions = c
+            i = Image.fromarray(np.zeros((size, size), dtype=np.uint8))
+            draw = ImageDraw.Draw(i)
+            width = random.randint(int(size * min_max_width[0]), int(size * min_max_width[1]))
+            offset = random.randint(int(-size * max_offset), int(size * max_offset)), random.randint(int(- size * max_offset), int(size * max_offset))
+            for sp, length, radians in zip(stroke_start_points, len_strokes, stroke_directions):
+                sp = (sp[0] + offset[0], sp[1] + offset[1])
+                x_vel = math.cos(radians) * length + random.randint(int(-size * max_target_offset), int(size * max_target_offset))
+                y_vel = math.sin(radians) * length + random.randint(int(-size * max_target_offset), int(size * max_target_offset))
+                new_p = (sp[0] + x_vel, sp[1] + y_vel)
+                stroke_directions.append(radians)
+                draw.line([round(x) for x in sp + new_p], fill=128, width=width)
+            a_i = np.array(i)
+            a_i[a_i == 128] = np.random.randint(200, 255, size=a_i.shape)[a_i == 128]
+            return Image.fromarray(a_i).filter(ImageFilter.GaussianBlur(.2))
+
+        generator_functions.append(g)
+    return generator_functions
+
+
+# g1,g2 = mnist_prior(2)
+
+# for i in [g1() for _ in range(10)]:
+#    display(i.resize((200,200)))
+
+from torchvision.transforms import ToTensor, ToPILImage
+
+
+def normalize(x):
+    return (x-x.mean())/(x.std()+.000001)
+
+from os import path, listdir
+import random
+
+def get_batch(batch_size, seq_len, num_features=None, noisy_std=None, only_train_for_last_idx=False, normalize_x=False, num_outputs=2, use_saved_from=None, **kwargs):  # num_features = 28*28=784
+    if use_saved_from is not None:
+        directory = path.join(use_saved_from, f'len_{seq_len}_out_{num_outputs}_features_{num_features}_bs_{batch_size}')
+        filename = random.choice(listdir(directory))
+        return torch.load(path.join(directory,filename))
+
+    size = math.isqrt(num_features)
+    assert size * size == num_features, 'num_features needs to be the square of an integer.'
+    if only_train_for_last_idx:
+        assert (seq_len-1) % num_outputs == 0
+
+    # assert seq_len % 2 == 0, "assert seq_len % 2 == 0"
+    batch = []
+    y = []
+    target_y = []
+    for b_i in range(batch_size):
+        gs = mnist_prior(num_outputs, size, **kwargs)
+        if only_train_for_last_idx:
+            generators = [i for i in range(len(gs)) for _ in range((seq_len-1) // num_outputs)]
+            random.shuffle(generators)
+            generators += [random.randint(0, len(gs) - 1)]
+            target = [-100 for _ in generators]
+            target[-1] = generators[-1]
+        else:
+            generators = [random.randint(0, len(gs) - 1) for _ in range(seq_len)]
+            target = generators
+        normalize_or_not = lambda x: normalize(x) if normalize_x else x
+        s = torch.cat([normalize_or_not(ToTensor()(gs[f_i]())) for f_i in generators], 0)
+        batch.append(s)
+        y.append(torch.tensor(generators))
+        target_y.append(torch.tensor(target))
+    x = torch.stack(batch, 1).view(seq_len, batch_size, -1)
+    y = torch.stack(y, 1)
+    target_y = torch.stack(target_y, 1)
+    return x,y,target_y
+
+DataLoader = get_batch_to_dataloader(get_batch)
+DataLoader.num_outputs = 2
+
+if __name__ == '__main__':
+    g1, g2 = mnist_prior(2, size=3)
+
+    # for i in range(10):
+    # print(PILToTensor()(g1()))
+    # display(ToPILImage()(PILToTensor()(g1())).resize((200,200)))
+    # display(g2().resize((200,200)))
+
+    size = 10
+    x, y = get_batch(1, 10, num_features=size * size)
+
+    x_ = x[..., :-1].squeeze(1)
+    last_y = x[..., -1].squeeze(1)
+    y = y.squeeze(1)
+
+    # print(y)
+
+    for i, y_, last_y_, x__ in zip(x_, y, last_y, x.squeeze(1)):
+        # print(y_)
+        # print(i.shape)
+        # print(x__)
+        img = ToPILImage()(i.view(size, size))
+        # display(img.resize((200,200)))
+
+    print(y, last_y)

prior-fitting/priors/utils.py

@@ -0,0 +1,102 @@
+import random
+
+import torch
+
+from utils import set_locals_in_self
+from itertools import repeat
+from .prior import PriorDataLoader
+from torch import nn
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.gridspec as gridspec
+import scipy.stats as stats
+
+def get_batch_to_dataloader(get_batch_method_):
+    class DL(PriorDataLoader):
+        get_batch_method = get_batch_method_
+
+        # Caution, you might need to set self.num_features manually if it is not part of the args.
+        def __init__(self, num_steps, fuse_x_y=False, **get_batch_kwargs):
+            set_locals_in_self(locals())
+            # The stuff outside the or is set as class attribute before instantiation.
+            self.num_features = get_batch_kwargs.get('num_features') or self.num_features
+            self.num_outputs = get_batch_kwargs.get('num_outputs') or self.num_outputs
+            print('DataLoader.__dict__', self.__dict__)
+
+        @staticmethod
+        def gbm(*args, fuse_x_y=True, **kwargs):
+            x, y, target_y = get_batch_method_(*args, **kwargs)
+            if fuse_x_y:
+                return torch.cat([x, torch.cat([torch.zeros_like(y[:1]), y[:-1]], 0).unsqueeze(-1).float()],
+                                 -1), target_y
+            else:
+                return (x, y), target_y
+
+        def __len__(self):
+            return self.num_steps
+
+        def __iter__(self):
+            return iter(self.gbm(**self.get_batch_kwargs, fuse_x_y=self.fuse_x_y) for _ in range(self.num_steps))
+
+
+    return DL
+
+
+def plot_features(data, targets):
+    if torch.is_tensor(data):
+        data = data.detach().cpu().numpy()
+        targets = targets.detach().cpu().numpy()
+    fig2 = plt.figure(constrained_layout=True, figsize=(12, 12))
+    spec2 = gridspec.GridSpec(ncols=data.shape[1], nrows=data.shape[1], figure=fig2)
+    for d in range(0, data.shape[1]):
+        for d2 in range(0, data.shape[1]):
+            sub_ax = fig2.add_subplot(spec2[d, d2])
+            sub_ax.scatter(data[:, d], data[:, d2],
+                           c=targets[:])
+
+
+def plot_prior(prior):
+    s = np.array([prior() for _ in range(0, 10000)])
+    count, bins, ignored = plt.hist(s, 50, density=True)
+    print(s.min())
+    plt.show()
+
+trunc_norm_sampler_f = lambda mu, sigma : lambda: stats.truncnorm((0 - mu) / sigma, (1 - mu) / sigma, loc=mu, scale=sigma).rvs(1)[0]
+beta_sampler_f = lambda a, b : lambda : np.random.beta(a, b)
+gamma_sampler_f = lambda a, b : lambda : np.random.gamma(a, b)
+uniform_sampler_f = lambda a, b : lambda : np.random.uniform(a, b)
+uniform_int_sampler_f = lambda a, b : lambda : np.random.randint(a, b)
+zipf_sampler_f = lambda a, b, c : lambda : min(b + np.random.zipf(a), c)
+scaled_beta_sampler_f = lambda a, b, scale, minimum : lambda : minimum + round(beta_sampler_f(a, b)() * (scale - minimum + 1) - 0.5)
+
+
+def normalize_data(data):
+    mean = data.mean(0)
+    std = data.std(0) + .000001
+    data = (data - mean) / std
+
+    return data
+
+
+def normalize_by_used_features_f(x, num_features_used, num_features):
+    return x / (num_features_used / num_features)
+
+
+class Binarize(nn.Module):
+    def __init__(self, p=0.5):
+        super().__init__()
+        self.p = p
+
+    def forward(self, x):
+        return (x > torch.median(x)).float()
+
+
+def order_by_y(x, y):
+    order = torch.argsort(y if random.randint(0, 1) else -y, dim=0)[:, 0, 0]
+    order = order.reshape(2, -1).transpose(0, 1).reshape(-1)#.reshape(seq_len)
+    x = x[order]  # .reshape(2, -1).transpose(0, 1).reshape(-1).flip([0]).reshape(seq_len, 1, -1)
+    y = y[order]  # .reshape(2, -1).transpose(0, 1).reshape(-1).reshape(seq_len, 1, -1)
+
+    return x, y
+
+

prior-fitting/requirements.txt

@@ -0,0 +1,13 @@
+# Recommend to use python >= 3.9
+gpytorch==1.5.0
+pyro-ppl==1.7.0
+torch==1.9.0
+scikit-learn==0.24.2
+pyyaml==5.4.1
+blitz-bayesian-pytorch==0.2.7
+seaborn==0.11.2
+xgboost==1.4.0
+tqdm==4.62.1
+numpy==1.21.2
+openml==0.12.2
+catboost==0.26.1

prior-fitting/tabular.py

@@ -0,0 +1,725 @@
+from catboost import CatBoostClassifier, Pool
+from sklearn.model_selection import GridSearchCV
+from sklearn.model_selection import KFold
+from sklearn.model_selection import ParameterGrid
+
+import pyro
+import pyro.distributions as dist
+from pyro.nn import PyroModule, PyroSample
+from pyro.infer.autoguide import AutoDiagonalNormal
+from pyro.infer import SVI, Trace_ELBO, Predictive, MCMC, NUTS
+from pytorch_tabnet.tab_model import TabNetClassifier, TabNetRegressor
+from sklearn.metrics import accuracy_score, roc_auc_score
+import argparse
+import itertools
+
+from train import train, get_weighted_single_eval_pos_sampler, Losses
+import priors
+import encoders
+from sklearn import preprocessing
+
+from sklearn.base import BaseEstimator, ClassifierMixin
+
+from torch import nn
+
+from datasets import *
+import xgboost as xgb
+import matplotlib.pyplot as plt
+import numpy as np
+
+import torch
+from sklearn import neighbors, datasets
+from sklearn.gaussian_process import GaussianProcessClassifier
+from sklearn.gaussian_process.kernels import RBF
+
+from priors.utils import trunc_norm_sampler_f, beta_sampler_f, gamma_sampler_f, uniform_sampler_f, zipf_sampler_f, scaled_beta_sampler_f, uniform_int_sampler_f
+
+from tqdm import tqdm
+import time
+import random
+
+import os
+
+CV = 5
+param_grid = {}
+metric_used = roc_auc_score
+
+def get_uniform_single_eval_pos_sampler(max_len):
+    """
+    Just sample any evaluation position with the same weight
+    :return: Sampler that can be fed to `train()` as `single_eval_pos_gen`.
+    """
+    return lambda: random.choices(range(max_len))[0]
+
+
+def get_mlp_prior_hyperparameters(config):
+    sigma_sampler = gamma_sampler_f(config["prior_sigma_gamma_k"], config["prior_sigma_gamma_theta"])
+    noise_std_sampler = gamma_sampler_f(config["prior_noise_std_gamma_k"], config["prior_noise_std_gamma_theta"])
+
+    mlp_prior_hyperparameters = (list(config["prior_nlayers_sampler"].values())[0]
+                                 , list(config["prior_emsize_sampler"].values())[0]
+                                 , config["prior_activations"]
+                                 , sigma_sampler
+                                 , noise_std_sampler
+                                 , list(config["prior_dropout_sampler"].values())[0]
+                                 , True
+                                 , list(config["prior_num_features_used_sampler"].values())[0]
+                                 , list(config["prior_causes_sampler"].values())[0] if config['prior_is_causal'] else None
+                                 , config["prior_is_causal"]
+                                 , config["prior_pre_sample_causes"] if config['prior_is_causal'] else None
+                                 , config["prior_pre_sample_weights"] if config['prior_is_causal'] else None
+                                 , config["prior_y_is_effect"] if config['prior_is_causal'] else None
+                                 , config["prior_order_y"]
+                                 , config["prior_normalize_by_used_features"]
+                                 , list(config["prior_categorical_feats"].values())[0] if config['prior_is_causal'] else None
+                                 , 0.0
+                                 )
+
+    return mlp_prior_hyperparameters
+
+
+def get_gp_mix_prior_hyperparameters(config):
+    return {'lengthscale_concentration': config["prior_lengthscale_concentration"],
+            'nu': config["prior_nu"],
+            'outputscale_concentration': config["prior_outputscale_concentration"],
+            'categorical_data': config["prior_y_minmax_norm"],
+            'y_minmax_norm': config["prior_lengthscale_concentration"],
+            'noise_concentration': config["prior_noise_concentration"],
+            'noise_rate': config["prior_noise_rate"]}
+
+
+def get_gp_prior_hyperparameters(config):
+
+
+    return (config['prior_noise']
+            , lambda : config['prior_outputscale']
+            , lambda : config['prior_lengthscale']  # lengthscale, Höher mehr sep
+            , True
+            , list(config['prior_num_features_used_sampler'].values())[0]
+            , config['prior_normalize_by_used_features']
+            , config['prior_order_y'])
+
+
+def get_meta_gp_prior_hyperparameters(config):
+    lengthscale_sampler = trunc_norm_sampler_f(config["prior_lengthscale_mean"], config["prior_lengthscale_mean"] * config["prior_lengthscale_std_f"])
+    outputscale_sampler = trunc_norm_sampler_f(config["prior_outputscale_mean"], config["prior_outputscale_mean"] * config["prior_outputscale_std_f"])
+
+    return (config['prior_noise']
+            , outputscale_sampler
+            , lengthscale_sampler  # lengthscale, Höher mehr sep
+            , True
+            , list(config['prior_num_features_used_sampler'].values())[0]
+            , config['prior_normalize_by_used_features']
+            , config['prior_order_y'])
+
+
+
+def get_model(config, device, eval_positions, should_train=True, verbose=False):
+    extra_kwargs = {}
+    if config['prior_type'] == 'mlp':
+        prior_hyperparameters = get_mlp_prior_hyperparameters(config)
+        model_proto = priors.mlp.DataLoader
+        extra_kwargs['batch_size_per_gp_sample'] = 8
+    elif config['prior_type'] == 'gp':
+        prior_hyperparameters = get_gp_prior_hyperparameters(config)
+        model_proto = priors.fast_gp.DataLoader
+    elif config['prior_type'] == 'custom_gp_mix':
+        prior_hyperparameters = get_meta_gp_prior_hyperparameters(config)
+        model_proto = priors.fast_gp.DataLoader
+    elif config['prior_type'] == 'gp_mix':
+        prior_hyperparameters = get_gp_mix_prior_hyperparameters(config)
+        model_proto = priors.fast_gp_mix.DataLoader
+    else:
+        raise Exception()
+
+    epochs = 0 if not should_train else config['epochs']
+    model = train(model_proto
+                  , Losses.bce
+                  , encoders.Linear
+                  , emsize=config['emsize']
+                  , nhead=config['nhead']
+                  , y_encoder_generator=encoders.Linear
+                  , pos_encoder_generator=None
+                  , batch_size=config['batch_size']
+                  , nlayers=config['nlayers']
+                  , nhid=config['emsize'] * config['nhid_factor']
+                  , epochs=epochs
+                  , warmup_epochs=epochs // 4
+                  , bptt=config['bptt']
+                  , gpu_device=device
+                  , dropout=config['dropout']
+                  , steps_per_epoch=100
+                  , single_eval_pos_gen=get_uniform_single_eval_pos_sampler(max(eval_positions) + 1)
+                  , extra_prior_kwargs_dict={
+            'num_features': config['num_features']
+            # , 'canonical_args': None
+            , 'fuse_x_y': False
+            , 'hyperparameters': prior_hyperparameters
+            , **extra_kwargs
+        }
+                  , lr=config['lr']
+                  , verbose=verbose)
+
+    return model
+
+
+## General eval
+
+def evaluate(datasets, model, method, bptt, eval_position_range, device, max_features=0, plot=False, extend_features=False, save=True, rescale_features=False, overwrite=False,
+             max_samples=40, path_interfix=''):
+    # eval_position_range: last entry is the one used to calculate metricuracy; up to index is used for training
+    result = {'metric': 'auc'}
+
+    metric_sum = 0
+    for [name, X, y, categorical_feats] in datasets:
+        result_ds = {}
+        path = f'/home/hollmann/prior-fitting/results/tabular/{path_interfix}/results_{method}_{name}.npy'
+        if (os.path.isfile(path)) and not overwrite:
+            with open(path, 'rb') as f:
+                result_ds = np.load(f, allow_pickle=True).tolist()
+                if 'time' in result_ds:
+                    result_ds[name+'_time'] = result_ds['time']
+                    del result_ds['time']
+                result.update(result_ds)
+                mean_metric = float(result[name + '_mean_metric_at_' + str(eval_position_range[-1])])
+                metric_sum += mean_metric
+                print(f'Loaded saved result for {name} (mean metric {mean_metric})')
+                continue
+
+        print('Evaluating ' + str(name))
+        rescale_features_factor = X.shape[1] / max_features if rescale_features and extend_features else 1.0
+        if extend_features:
+            X = torch.cat([X, torch.zeros((X.shape[0], max_features - X.shape[1]))], -1)
+
+        start_time = time.time()
+        ds_result = evaluate_dataset(X.to(device), y.to(device), categorical_feats, model, bptt, eval_position_range,
+                               rescale_features=rescale_features_factor, max_samples=max_samples)
+        elapsed = time.time() - start_time
+
+        for i, r in enumerate(ds_result):
+            metric, outputs, ys = r
+            if save:
+                result_ds[name + '_per_ds_metric_at_' + str(eval_position_range[i])] = metric
+                result_ds[name + '_outputs_at_' + str(eval_position_range[i])] = outputs
+                result_ds[name + '_ys_at_' + str(eval_position_range[i])] = ys
+
+            result_ds[name + '_mean_metric_at_' + str(eval_position_range[i])] = metric_used(ys.detach().cpu().flatten(), outputs.flatten())
+            result_ds[name + '_time'] = elapsed
+
+        if save:
+            with open(path, 'wb') as f:
+                np.save(f, result_ds)
+
+        result.update(result_ds)
+        metric_sum += float(metric[-1].mean())
+
+    for pos in eval_position_range:
+        result[f'mean_metric_at_{pos}'] = np.array([result[d[0] + '_mean_metric_at_' + str(pos)] for d in datasets]).mean()
+
+    result['mean_metric'] = np.array([result['mean_metric_at_' + str(pos)] for pos in eval_position_range]).mean()
+
+    return result
+
+
+def evaluate_dataset(X, y, categorical_feats, model, bptt, eval_position_range, plot=False, rescale_features=1.0,
+                     max_samples=40):
+    result = []
+    for eval_position in eval_position_range:
+        r = evaluate_position(X, y, categorical_feats, model, bptt, eval_position, rescale_features=rescale_features,
+                              max_samples=max_samples)
+        result.append(r)
+        print('\t Eval position ' + str(eval_position) + ' done..')
+
+    if plot:
+        plt.plot(np.array(list(eval_position_range)), np.array([r.mean() for r in result]))
+
+    return result
+
+
+def evaluate_position(X, y, categorical_feats, model, bptt, eval_position, rescale_features=1.0, max_samples=40):
+    # right now permutation style is to test performance on one before the last element
+    # eval_position = bptt - eval_positions
+
+    # TODO: Make sure that no bias exists
+    # assert(eval_position % 2 == 0)
+
+    eval_xs = []
+    eval_ys = []
+    num_evals = len(X) - bptt  # len(X)-bptt-(bptt-eval_position)+1
+
+    # Generate permutations of evaluation data
+    #     with torch.random.fork_rng():
+    #         torch.random.manual_seed(13)
+    #         ps = [torch.randperm(2*(bptt - eval_position)) for _ in range(num_evals)]
+
+    for i in range(num_evals):
+        # Select chunk of data with extra evaluation positions that can be discarded
+        #         x_ = X[i:i+bptt+(bptt-eval_position)].clone()
+        #         y_ = y[i:i+bptt+(bptt-eval_position)].clone()
+
+        #         # Permutate evaluation positions
+        #         perm_range = slice(eval_position,bptt+(bptt - eval_position))
+        #         x_[perm_range] = x_[perm_range][ps[i]]
+        #         y_[perm_range] = y_[perm_range][ps[i]]
+
+        #         # Discard extra evaluation positions
+        #         x_ = x_[0:bptt]
+        #         y_ = y_[0:bptt]
+
+        x_ = X[i:i + bptt].clone()
+        y_ = y[i:i + bptt].clone()
+
+        eval_xs.append(x_)
+        eval_ys.append(y_)
+
+    # eval data will be ordered in training range and
+    #   will be a random subset of 2*eval_position data points in eval positions
+    eval_xs = torch.stack(eval_xs, 1)
+    eval_ys = torch.stack(eval_ys, 1)
+
+    # Limit to N samples per dataset
+    with torch.random.fork_rng():
+        torch.random.manual_seed(13)
+        sel = torch.randperm(eval_xs.shape[1])
+        eval_xs = eval_xs[:, sel[0:max_samples], :]
+        eval_ys = eval_ys[:, sel[0:max_samples]]
+    #
+    # if quantile_transform:
+    #     for sample in range(0, eval_xs.shape[1]):
+    #         quantile_transformer = preprocessing.QuantileTransformer(random_state=0, n_quantiles=eval_xs.shape[0])
+    #         quantile_transformer.fit(eval_xs[:eval_position, sample].cpu())
+    #         eval_xs[:, sample] = torch.tensor(quantile_transformer.transform(eval_xs[:, sample].cpu()))
+
+    if isinstance(model, nn.Module):
+        model.eval()
+        outputs = np.zeros(shape=(len(list(range(eval_position, eval_xs.shape[0]))), eval_xs.shape[1]))
+        for i, pos in enumerate(range(eval_position, eval_xs.shape[0])):
+            eval_x = torch.cat([eval_xs[:eval_position], eval_xs[pos].unsqueeze(0)])
+            eval_y = eval_ys[:eval_position]
+
+            # Center data using training positions so that it matches priors
+            mean = eval_x.mean(0)
+            std = eval_x.std(0) + .000001
+            eval_x = (eval_x - mean) / std
+            eval_x = eval_x / rescale_features
+
+            output = torch.sigmoid(model((eval_x, eval_y.float()), single_eval_pos=eval_position)).squeeze(-1)
+            outputs[i, :] = output.detach().cpu().numpy()
+
+        metric_per_t = np.array([metric_used(eval_ys[eval_position:, i].cpu(), outputs[:, i]) for i in range(eval_xs.shape[1])])
+        return metric_per_t, outputs, eval_ys[eval_position:]
+    else:
+        metric_eval_pos, outputs = batch_pred(model, eval_xs, eval_ys, categorical_feats, start=eval_position)
+
+        return metric_eval_pos, outputs, eval_ys[eval_position:]
+
+
+def batch_pred(metric_function, eval_xs, eval_ys, categorical_feats, start=2):
+    metrics = []
+    outputs = []
+    # for i in tqdm(list(range(start,len(eval_xs)))):
+    eval_splits = list(zip(eval_xs.transpose(0, 1), eval_ys.transpose(0, 1)))
+    for eval_x, eval_y in tqdm(eval_splits):  # eval x is One sample i.e. bptt x num_features
+        mean = eval_x[:start].mean(0)
+        std = eval_x[:start].std(0) + .000001
+        eval_x = (eval_x - mean) / std
+
+        metric, output = metric_function(eval_x[:start], eval_y[:start], eval_x[start:], eval_y[start:], categorical_feats)
+        metrics += [metric]
+        outputs += [output]
+    #     metrics_per_t.append(metric_sum/eval_xs.shape[1])
+    return np.array(metrics), np.array(outputs).T
+
+## Ridge
+
+
+from sklearn.linear_model import RidgeClassifier
+# param_grid['ridge'] = {'alpha': [0, 0.1, .5, 1.0, 2.0], 'fit_intercept': [True, False]} # 'normalize': [False],
+def ridge_metric(x, y, test_x, test_y, cat_features):
+    import warnings
+    def warn(*args, **kwargs):
+        pass
+
+    warnings.warn = warn
+
+    x, y, test_x, test_y = x.cpu(), y.cpu(), test_x.cpu(), test_y.cpu()
+
+    clf = RidgeClassifier()
+
+    # create a dictionary of all values we want to test for n_neighbors
+    # use gridsearch to test all values for n_neighbors
+    clf = GridSearchCV(clf, param_grid['ridge'], cv=min(CV, x.shape[0]//2))
+    # fit model to data
+    clf.fit(x, y.long())
+
+    pred = clf.decision_function(test_x)
+    metric = metric_used(test_y.cpu().numpy(), pred)
+
+    return metric, pred
+
+
+from sklearn.linear_model import LogisticRegression
+param_grid['logistic'] = {'solver': ['saga'], 'penalty': ['l1', 'l2', 'none'], 'tol': [1e-2, 1e-4, 1e-10], 'max_iter': [500], 'fit_intercept': [True, False], 'C': [1e-5, 0.001, 0.01, 0.1, 1.0, 2.0]} # 'normalize': [False],
+def logistic_metric(x, y, test_x, test_y, cat_features):
+    import warnings
+    def warn(*args, **kwargs):
+        pass
+
+    warnings.warn = warn
+
+    x, y, test_x, test_y = x.cpu(), y.cpu(), test_x.cpu(), test_y.cpu()
+
+    clf = LogisticRegression()
+
+    # create a dictionary of all values we want to test for n_neighbors
+    # use gridsearch to test all values for n_neighbors
+    clf = GridSearchCV(clf, param_grid['logistic'], cv=min(CV, x.shape[0]//2))
+    # fit model to data
+    clf.fit(x, y.long())
+
+    pred = clf.predict_proba(test_x)[:, 1]
+    metric = metric_used(test_y.cpu().numpy(), pred)
+
+    return metric, pred
+
+
+## KNN
+param_grid['knn'] = {'n_neighbors (max number of samples)': np.arange(1, 6)}
+def knn_metric(x, y, test_x, test_y, cat_features):
+    x, y, test_x, test_y = x.cpu(), y.cpu(), test_x.cpu(), test_y.cpu()
+
+    clf = neighbors.KNeighborsClassifier()  # min(param['n_neighbors'],len(y)))
+    param_grid_knn = {'n_neighbors': np.arange(1, min(6, len(y) - 1))}
+    # create a dictionary of all values we want to test for n_neighbors
+    # use gridsearch to test all values for n_neighbors
+    clf = GridSearchCV(clf, param_grid_knn, cv=min(CV, x.shape[0]//2))
+    # fit model to data
+    clf.fit(x, y.long())
+
+    # print(clf.best_params_)
+
+    # clf.fit(x, y.long())
+    pred = clf.predict_proba(test_x)[:, 1]
+
+    metric = metric_used(test_y.cpu().numpy(), pred)
+
+    return metric, pred
+
+
+## Bayesian NN
+class BayesianModel(PyroModule):
+    def __init__(self, model_spec, device='cuda'):
+        super().__init__()
+
+        self.device = device
+        self.num_features = model_spec['num_features']
+
+        mu, sigma = torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)
+
+        self.fc1 = PyroModule[nn.Linear](self.num_features, model_spec['embed'])
+        self.fc1.weight = PyroSample(
+            dist.Normal(mu, sigma).expand([model_spec['embed'], self.num_features]).to_event(2))
+        self.fc1.bias = PyroSample(dist.Normal(mu, sigma).expand([model_spec['embed']]).to_event(1))
+        self.fc2 = PyroModule[nn.Linear](model_spec['embed'], 2)
+        self.fc2.weight = PyroSample(dist.Normal(mu, sigma).expand([2, model_spec['embed']]).to_event(2))
+        self.fc2.bias = PyroSample(dist.Normal(mu, sigma).expand([2]).to_event(1))
+
+        self.model = torch.nn.Sequential(self.fc1, self.fc2)
+
+        self.to(self.device)
+
+    def forward(self, x=None, y=None, seq_len=1):
+        if x is None:
+            with pyro.plate("x_plate", seq_len):
+                d_ = dist.Normal(torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)).expand(
+                    [self.num_features]).to_event(1)
+                x = pyro.sample("x", d_)
+
+        out = self.model(x)
+        mu = out.squeeze()
+        softmax = torch.nn.Softmax(dim=1)
+        # sigma = pyro.sample("sigma", dist.Uniform(torch.tensor([0.0]).to(self.device), torch.tensor([1.0]).to(self.device)))
+        with pyro.plate("data", out.shape[0]):
+            # d_ = dist.Normal(mu, sigma)
+            # obs = pyro.sample("obs", d_, obs=y)
+            s = softmax(mu)
+            obs = pyro.sample('obs', dist.Categorical(probs=s), obs=y).float()
+
+        return x, obs
+
+
+class BayesianNNClassifier(BaseEstimator, ClassifierMixin):
+
+    def __init__(self, num_features, n_layers, embed, lr, device):
+        self.num_pred_samples = 400
+        self.num_steps = 400
+        self.embed = embed
+        self.n_layers = n_layers
+        self.lr = lr
+        self.num_features = num_features
+        self.device = device
+
+    def fit(self, X, y):
+        model_spec = {'nlayers': 2, 'embed': self.embed, 'num_features': self.num_features}
+
+        self.model = BayesianModel(model_spec, device=self.device)
+        self.guide = AutoDiagonalNormal(self.model).to(self.device)
+        self.adam = pyro.optim.Adam({"lr": self.lr})
+        self.svi = SVI(self.model, self.guide, self.adam, loss=Trace_ELBO())
+
+        pyro.clear_param_store()
+
+        X = X.to(self.device)
+        y = y.to(self.device)
+
+        for epoch in tqdm(range(0, self.num_steps)):
+            loss = self.svi.step(X, y)
+
+        # Return the classifier
+        return self
+
+    def predict(self, X):
+        X = X.to(self.device)
+        predictive = Predictive(self.model, guide=self.guide, num_samples=self.num_pred_samples)
+        preds = predictive(X)['obs']
+        preds_means = preds.float().mean(axis=0).detach().cpu()
+        preds_hard = preds_means > 0.5
+
+        return preds_hard.long()
+
+    def predict_proba(self, X):
+        X = X.to(self.device)
+        predictive = Predictive(self.model, guide=self.guide, num_samples=self.num_pred_samples)
+        preds = predictive(X)['obs']
+        preds_means = preds.float().mean(axis=0).detach().cpu()
+
+        return preds_means
+
+    def score(self, X, y):
+        return super().score(X, y)
+
+param_grid['bayes'] = {'embed': [5, 10, 30, 64], 'lr': [1e-3, 1e-4], 'num_training_steps': [400], 'num_samples_for_prediction': [400]}
+def bayes_net_metric(x, y, test_x, test_y, cat_features):
+    device = x.device
+
+    clf = BayesianNNClassifier(x.shape[1], 2, 1, 1e-3, device)
+    # create a dictionary of all values we want to test for n_neighbors
+    # use gridsearch to test all values for n_neighbors
+    clf = GridSearchCV(clf, param_grid['bayes'], cv=5)
+    # fit model to data
+    clf.fit(x.cpu(), y.long().cpu())
+
+    pred = clf.predict_proba(test_x)
+    metric = metric_used(test_y.cpu().numpy(), pred.cpu().numpy())
+
+    return metric, pred
+
+## GP
+param_grid['gp'] = {'params_y_scale': [0.05, 0.1, 0.5, 1.0, 5.0, 10.0],
+    'params_length_scale': [0.1, 0.5, 1.0, 2.0]}
+def gp_metric(x, y, test_x, test_y, cat_features):
+    import warnings
+    def warn(*args, **kwargs):
+        pass
+    warnings.warn = warn
+
+    x, y, test_x, test_y = x.cpu(), y.cpu(), test_x.cpu(), test_y.cpu()
+
+    clf = GaussianProcessClassifier()
+    # create a dictionary of all values we want to test for n_neighbors
+    params_y_scale = [0.05, 0.1, 0.5, 1.0, 5.0, 10.0]# 0.000001, 0.00001,
+    params_length_scale = [0.1, 0.5, 1.0, 2.0] # 0.01,
+    param_grid = {'kernel': [y * RBF(l) for (y, l) in list(itertools.product(params_y_scale, params_length_scale))]}
+    # use gridsearch to test all values for n_neighbors
+    clf = GridSearchCV(clf, param_grid, cv=min(CV, x.shape[0]//2))
+    # fit model to data
+    clf.fit(x, y.long())
+    pred = clf.predict_proba(test_x)[:, 1]
+    metric = metric_used(test_y.cpu().numpy(), pred)
+
+    return metric, pred
+
+
+## Tabnet
+# https://github.com/dreamquark-ai/tabnet
+param_grid['tabnet'] = {'n_d': [2, 4], 'n_steps': [2,4,6], 'gamma': [1.3], 'optimizer_params': [{'lr': 2e-2}, {'lr': 2e-1}]}
+#param_grid['tabnet'] = {'n_d': [2], 'n_steps': [2], 'optimizer_params': [{'lr': 2e-2}, {'lr': 2e-1}]}
+def tabnet_metric(x, y, test_x, test_y, cat_features):
+    x, y, test_x, test_y = x.cpu().numpy(), y.cpu().numpy(), test_x.cpu().numpy(), test_y.cpu().numpy()
+
+    mean_metrics = []
+    mean_best_epochs = []
+
+    for params in list(ParameterGrid(param_grid['tabnet'])):
+        kf = KFold(n_splits=min(5, x.shape[0]//2), random_state=None, shuffle=False)
+        metrics = []
+        best_epochs = []
+        for train_index, test_index in kf.split(x):
+            X_train, X_valid, y_train, y_valid = x[train_index], x[test_index], y[train_index], y[test_index]
+
+            clf = TabNetClassifier(verbose=True, cat_idxs=cat_features, n_a=params['n_d'], **params)
+
+            clf.fit(
+                X_train, y_train,
+                #eval_set=[(X_valid, y_valid)], patience=15
+            )
+
+            metric = metric_used(test_y.cpu().numpy(), clf.predict(X_valid))
+            metrics += [metric]
+            #best_epochs += [clf.best_epoch]
+        mean_metrics += [np.array(metrics).mean()]
+        #mean_best_epochs += [np.array(best_epochs).mean().astype(int)]
+
+    mean_metrics = np.array(mean_metrics)
+    #mean_best_epochs = np.array(mean_best_epochs)
+    params_used = np.array(list(ParameterGrid(param_grid['tabnet'])))
+
+    best_idx = np.argmax(mean_metrics)
+    #print(params_used[best_idx])
+    clf = TabNetClassifier(cat_idxs=cat_features, **params_used[best_idx])
+
+    clf.fit(
+        x, y#, max_epochs=mean_best_epochs[best_idx]
+    )
+
+    pred = 1 - clf.predict_proba(test_x)[:,0]
+    metric = metric_used(test_y, pred)
+
+    #print(metric, clf.predict(test_x), pred)
+
+    return metric, pred
+
+
+# Catboost
+param_grid['catboost'] = {'learning_rate': [0.1, 0.5, 1.0],
+            'depth': [2, 4, 7],
+            'l2_leaf_reg': [0.0, 0.5, 1],
+            'iterations': [10, 40, 70],
+                          'loss_function': ['Logloss']}
+def catboost_metric(x, y, test_x, test_y, categorical_feats):
+    import warnings
+    def warn(*args, **kwargs):
+        pass
+
+    warnings.warn = warn
+
+    x, y, test_x, test_y = x.numpy(), y.numpy(), test_x.numpy(), test_y.numpy()
+
+    def make_pd_from_np(x):
+        data = pd.DataFrame(x)
+        for c in categorical_feats:
+            data.iloc[:, c] = data.iloc[:, c].astype('int')
+        return data
+
+    x = make_pd_from_np(x)
+    test_x = make_pd_from_np(test_x)
+
+    model = CatBoostClassifier(iterations=2,
+                               depth=2,
+                               learning_rate=1,
+                               loss_function='Logloss',
+                               logging_level='Silent')
+
+    grid_search_result = model.grid_search(param_grid['catboost'],
+                                           X=x,
+                                           y=y,
+                                           cv=5,
+                                           plot=False,
+                                           verbose=False)  # randomized_search with n_iter
+
+    # model.fit(x, y)
+    pred = model.predict_proba(test_x)[:, 1]
+    metric = metric_used(test_y.cpu().numpy(), pred)
+
+    return metric, pred
+
+
+# XGBoost
+param_grid['xgb'] = {
+        'min_child_weight': [0.5, 1.0],
+        'learning_rate': [0.02, 0.2],
+        #'gamma': [0.1, 0.2, 0.5, 1, 2],
+        'subsample': [0.5, 0.8],
+        'max_depth': [1, 2],
+        'colsample_bytree': [0.8], #0.5,
+        'eval_metric': ['logloss'],
+        'n_estimators': [100]
+    }
+def xgb_metric(x, y, test_x, test_y, cat_features):
+    x, y, test_x, test_y = x.numpy(), y.numpy().astype(int), test_x.numpy(), test_y.numpy().astype(int)
+
+    clf = xgb.XGBClassifier(use_label_encoder=False)
+
+    # {'num_round': [2,5,10,20], 'max_depth': [1, 2,4,6,8], 'eta': [.1, .01, .001], 'eval_metric': 'logloss'}
+    # use gridsearch to test all values for n_neighbors
+    clf = GridSearchCV(clf, param_grid['xgb'], cv=5, n_jobs=4, verbose=2)
+    # fit model to data
+    clf.fit(x, y.astype(int))
+
+    print(clf.best_params_)
+
+    # clf.fit(x, y.long())
+    pred = clf.predict_proba(test_x)[:, 1]
+    metrics = ((pred > 0.5) == test_y).astype(float).mean()
+    return metrics, pred
+
+def get_default_spec(test_datasets, valid_datasets):
+    bptt = 100
+    eval_positions = [30] #list(range(6, 42, 2))  # list(range(10, bptt-10, 20)) + [bptt-10]
+    max_features = max([X.shape[1] for (_, X, _, _) in test_datasets] + [X.shape[1] for (_, X, _, _) in valid_datasets])
+    max_samples = 20
+
+    return bptt, eval_positions, max_features, max_samples
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--method', default='ridge', type=str)
+    parser.add_argument('--did', default=-1, type=int)
+    parser.add_argument('--overwrite', default=False, type=bool)
+    args = parser.parse_args()
+
+    test_datasets, _ = load_openml_list(test_dids_classification)
+    valid_datasets, _ = load_openml_list(valid_dids_classification)
+
+    selector = 'test'
+    ds = valid_datasets if selector == 'valid' else test_datasets
+    if args.did > -1:
+        ds = ds[args.did:args.did+1]
+
+    bptt, eval_positions, max_features, max_samples = get_default_spec(test_datasets, valid_datasets)
+
+    if args.method == 'bayes':
+        clf = bayes_net_metric
+        device = 'cpu'
+    elif args.method == 'gp':
+        clf = gp_metric
+        device = 'cpu'
+    elif args.method == 'ridge':
+        clf = ridge_metric
+        device = 'cpu'
+    elif args.method == 'knn':
+        clf = knn_metric
+        device = 'cpu'
+    elif args.method == 'catboost':
+        clf = catboost_metric
+        device = 'cpu'
+    elif args.method == 'tabnet':
+        clf = tabnet_metric
+        device = 'cpu'
+    elif args.method == 'xgb':
+        # Uses lots of cpu so difficult to time
+        clf = xgb_metric
+        device = 'cpu'
+    elif args.method == 'logistic':
+        clf = logistic_metric
+        device = 'cpu'
+    else:
+        clf = None
+        device = 'cpu'
+
+    start_time = time.time()
+    result = evaluate(ds, clf, args.method, bptt, eval_positions, device=device, max_samples=max_samples, overwrite=args.overwrite, save=True)
+    result['time_spent'] = time.time() - start_time
+
+    with open(f'/home/hollmann/prior-fitting/results/tabular/results_{selector}_{args.method}.npy', 'wb') as f:
+        np.save(f, result)

prior-fitting/train.py

@@ -0,0 +1,288 @@
+import argparse
+import time
+import yaml
+
+import torch
+from torch import nn
+from transformer import TransformerModel
+from bar_distribution import BarDistribution, FullSupportBarDistribution, get_bucket_limits
+from utils import get_cosine_schedule_with_warmup, get_openai_lr, StoreDictKeyPair, get_weighted_single_eval_pos_sampler, get_uniform_single_eval_pos_sampler
+import priors
+import encoders
+import positional_encodings
+
+class Losses():
+    gaussian = nn.GaussianNLLLoss(full=True, reduction='none')
+    mse = nn.MSELoss(reduction='none')
+    ce = nn.CrossEntropyLoss(reduction='none')
+    bce = nn.BCEWithLogitsLoss(reduction='none')
+    get_BarDistribution = BarDistribution
+
+
+def train(priordataloader_class, criterion, encoder_generator, emsize=200, nhid=200, nlayers=6, nhead=2, dropout=0.2,
+          epochs=10, steps_per_epoch=100, batch_size=200, bptt=10, lr=None, warmup_epochs=10, input_normalization=False,
+          y_encoder_generator=None, pos_encoder_generator=None, decoder=None, extra_prior_kwargs_dict={}, scheduler=get_cosine_schedule_with_warmup,
+          load_weights_from_this_state_dict=None, validation_period=10, single_eval_pos_gen=None, gpu_device='cuda:0',
+          aggregate_k_gradients=1, verbose=True
+          ):
+
+    device = gpu_device if torch.cuda.is_available() else 'cpu:0'
+    print(f'Using {device} device')
+    dl = priordataloader_class(num_steps=steps_per_epoch, batch_size=batch_size, seq_len=bptt, **extra_prior_kwargs_dict)
+
+    encoder = encoder_generator(dl.num_features+1 if dl.fuse_x_y else dl.num_features,emsize)
+    n_out = dl.num_outputs
+    if isinstance(criterion, nn.GaussianNLLLoss):
+        n_out *= 2
+    elif isinstance(criterion, BarDistribution) or "BarDistribution" in criterion.__class__.__name__: # TODO remove this fix (only for dev)
+        assert n_out == 1
+        n_out = criterion.num_bars
+    model = TransformerModel(encoder, n_out, emsize, nhead, nhid, nlayers, dropout,
+                             y_encoder=y_encoder_generator(1, emsize), input_normalization=input_normalization,
+                             pos_encoder=(pos_encoder_generator or positional_encodings.NoPositionalEncoding)(emsize, bptt*2),
+                             decoder=decoder
+                             )
+    model.criterion = criterion
+    if load_weights_from_this_state_dict is not None:
+        model.load_state_dict(load_weights_from_this_state_dict)
+    model.to(device)
+
+
+    # learning rate
+    if lr is None:
+        lr = get_openai_lr(model)
+        print(f"Using OpenAI max lr of {lr}.")
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+    scheduler = scheduler(optimizer, warmup_epochs, epochs)
+
+    def train():
+        model.train()  # Turn on the train mode
+        total_loss = 0.
+        total_positional_losses = 0.
+        total_positional_losses_recorded = 0
+        start_time = time.time()
+        before_get_batch = time.time()
+        assert len(dl) % aggregate_k_gradients == 0, 'Please set the number of steps per epoch s.t. `aggregate_k_gradients` divides it.'
+        for batch, (data, targets) in enumerate(dl):
+            time_to_get_batch = time.time() - before_get_batch
+            before_forward = time.time()
+            single_eval_pos = single_eval_pos_gen() if callable(single_eval_pos_gen) else single_eval_pos_gen
+            output = model(tuple(e.to(device) for e in data) if isinstance(data, tuple) else data.to(device)
+                           , single_eval_pos=single_eval_pos)
+
+            forward_time = time.time() - before_forward
+
+            if single_eval_pos is not None:
+                targets = targets[single_eval_pos:]
+
+            if isinstance(criterion, nn.GaussianNLLLoss):
+                assert output.shape[-1] == 2, \
+                    'need to write a little bit of code to handle multiple regression targets at once'
+
+                mean_pred = output[..., 0]
+                var_pred = output[..., 1].abs()
+                losses = criterion(mean_pred.flatten(), targets.to(device).flatten(), var=var_pred.flatten())
+            elif isinstance(criterion, (nn.MSELoss, nn.BCEWithLogitsLoss)):
+                losses = criterion(output.flatten(), targets.to(device).flatten())
+            else:
+                losses = criterion(output.reshape(-1, n_out), targets.to(device).flatten())
+            losses = losses.view(*output.shape[0:2]).squeeze(-1)
+
+
+            loss = losses.mean()
+            loss.backward()
+            if batch % aggregate_k_gradients == aggregate_k_gradients - 1:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
+                optimizer.step()
+                optimizer.zero_grad()
+
+            step_time = time.time() - before_forward
+
+            total_loss += loss.item()
+            total_positional_losses += losses.mean(1).cpu().detach() if single_eval_pos is None else \
+                nn.functional.one_hot(torch.tensor(single_eval_pos), bptt)*loss.cpu().detach()
+
+            total_positional_losses_recorded += torch.ones(bptt) if single_eval_pos is None else \
+                nn.functional.one_hot(torch.tensor(single_eval_pos), bptt)
+
+            before_get_batch = time.time()
+        return total_loss / steps_per_epoch, (
+                    total_positional_losses / total_positional_losses_recorded).tolist(), time_to_get_batch, forward_time, step_time
+
+    best_val_loss = float("inf")
+    best_model = None
+    total_loss = float('inf')
+    total_positional_losses = float('inf')
+    for epoch in range(1, epochs + 1):
+        epoch_start_time = time.time()
+        total_loss, total_positional_losses, time_to_get_batch, forward_time, step_time = train()
+        if hasattr(dl, 'validate') and epoch % validation_period == 0:
+            with torch.no_grad():
+                val_score = dl.validate(model)
+        else:
+            val_score = None
+
+        if verbose:
+            print('-' * 89)
+            print(
+                f'| end of epoch {epoch:3d} | time: {(time.time() - epoch_start_time):5.2f}s | mean loss {total_loss:5.2f} | '
+                f"pos losses {','.join([f'{l:5.2f}' for l in total_positional_losses])}, lr {scheduler.get_last_lr()[0]}"
+                f' data time {time_to_get_batch:5.2f} step time {step_time:5.2f}'
+                f' forward time {forward_time:5.2f}' + (f'val score {val_score}' if val_score is not None else ''))
+            print('-' * 89)
+
+        scheduler.step()
+    return total_loss, total_positional_losses, model.to('cpu')
+
+def _parse_args(config_parser, parser):
+    # Do we have a config file to parse?
+    args_config, remaining = config_parser.parse_known_args()
+    if args_config.config:
+        with open(args_config.config, 'r') as f:
+            cfg = yaml.safe_load(f)
+            parser.set_defaults(**cfg)
+
+    # The main arg parser parses the rest of the args, the usual
+    # defaults will have been overridden if config file specified.
+    args = parser.parse_args(remaining)
+
+    # Cache the args as a text string to save them in the output dir later
+    args_text = yaml.safe_dump(args.__dict__, default_flow_style=False)
+    return args, args_text
+
+
+if __name__ == '__main__':
+    config_parser = argparse.ArgumentParser(description='Only used as a first parser for the config file path.')
+    config_parser.add_argument('--config')
+    parser = argparse.ArgumentParser()
+    parser.add_argument('prior')
+    parser.add_argument('--loss_function', default='barnll')
+    # Optional Arg's for `--loss_function barnll`
+    parser.add_argument('--min_y', type=float, help='barnll can only model y in strict ranges, this is the minimum y can take.')
+    parser.add_argument('--max_y', type=float, help='barnll can only model y in strict ranges, this is the maximum y can take.')
+    parser.add_argument('--num_buckets', default=100, type=int)
+    #parser.add_argument('--num_features', default=None, type=int, help='Specify depending on the prior.')
+    parser.add_argument("--extra_prior_kwargs_dict", default={'fuse_x_y': False}, dest="extra_prior_kwargs_dict", action=StoreDictKeyPair, nargs="+", metavar="KEY=VAL", help='Specify depending on the prior.')
+    parser.add_argument('--encoder', default='linear', type=str, help='Specify depending on the prior.')
+    parser.add_argument('--y_encoder', default='linear', type=str, help='Specify depending on the prior. You should specify this if you do not fuse x and y.')
+    parser.add_argument('--pos_encoder', default='sinus', type=str, help='Specify depending on the prior.')
+    parser.add_argument('--bptt', default=10, type=int)
+    parser.add_argument('--epochs', default=200, type=int)
+    parser.add_argument('--warmup_epochs', default=50, type=int)
+    parser.add_argument('--validation_period', default=10, type=int)
+    parser.add_argument('--permutation_invariant_max_eval_pos', default=None, type=int, help='Set this to an int to ')
+    parser.add_argument('--permutation_invariant_sampling', default='weighted', help="Only relevant if --permutation_invariant_max_eval_pos is set.")
+
+    # these can likely be mostly left at defaults
+    parser.add_argument('--emsize', default=512, type=int) # sometimes even larger is better e.g. 1024
+    parser.add_argument('--nlayers', default=6, type=int)
+    parser.add_argument('--nhid', default=None, type=int) # 2*emsize is the default
+    parser.add_argument('--nhead', default=4, type=int) # nhead = emsize / 64 in the original paper
+    parser.add_argument('--dropout', default=.0, type=float)
+    parser.add_argument('--steps_per_epoch', default=10, type=int)
+    parser.add_argument('--batch_size', default=1000, type=int)
+    parser.add_argument('--lr', '--learning_rate', default=.001, type=float) # try also .0003, .0001, go lower with lower batch size
+
+    args, _ = _parse_args(config_parser, parser)
+
+    if args.nhid is None:
+        args.nhid = 2*args.emsize
+
+    prior = args.__dict__.pop('prior')
+
+    if prior == 'gp':
+        prior = priors.fast_gp.DataLoader
+    elif prior == 'ridge':
+        prior = priors.ridge.DataLoader
+    elif prior == 'stroke':
+        prior = priors.stroke.DataLoader
+    elif prior == 'mix_gp':
+        prior = priors.fast_gp_mix.DataLoader
+    else:
+        raise NotImplementedError(f'Prior == {prior}.')
+
+
+    loss_function = args.__dict__.pop('loss_function')
+
+    criterion = nn.GaussianNLLLoss(reduction='none', full=True)
+    classificiation_criterion = nn.CrossEntropyLoss(reduction='none')
+    num_buckets = args.__dict__.pop('num_buckets')
+    max_y = args.__dict__.pop('max_y')
+    min_y = args.__dict__.pop('min_y')
+    # criterion = nn.MSELoss(reduction='none')
+
+    def get_y_sample():
+        dl = prior(num_steps=1, batch_size=args.batch_size * args.steps_per_epoch, seq_len=args.bptt,
+                   **args.extra_prior_kwargs_dict)
+        y_sample = next(iter(dl))[-1]
+        print(f'Creating Bar distribution with borders from y sample of size {y_sample.numel()}')
+        return y_sample
+
+    if loss_function == 'ce':
+        criterion = nn.CrossEntropyLoss(reduction='none')
+    elif loss_function == 'gaussnll':
+        criterion = nn.GaussianNLLLoss(reduction='none', full=True)
+    elif loss_function == 'mse':
+        criterion = nn.MSELoss(reduction='none')
+    elif loss_function == 'barnll':
+        criterion = BarDistribution(borders=get_bucket_limits(num_buckets, full_range=(min_y,max_y)))
+    elif loss_function == 'adaptivebarnll':
+        borders = get_bucket_limits(num_buckets, ys=get_y_sample(), full_range=(min_y,max_y))
+        criterion = BarDistribution(borders=borders)
+    elif loss_function == 'adaptivefullsupportbarnll':
+        assert min_y is None and max_y is None, "Please do not specify `min_y` and `max_y` with `unboundedadaptivebarnll`."
+        borders = get_bucket_limits(num_buckets, ys=get_y_sample())
+        criterion = FullSupportBarDistribution(borders=borders)
+    else:
+        raise NotImplementedError(f'loss_function == {loss_function}.')
+
+
+
+    encoder = args.__dict__.pop('encoder')
+    y_encoder = args.__dict__.pop('y_encoder')
+
+    def get_encoder_generator(encoder):
+        if encoder == 'linear':
+            encoder_generator = encoders.Linear
+        elif encoder == 'mlp':
+            encoder_generator = encoders.MLP
+        elif encoder == 'positional':
+            encoder_generator = encoders.Positional
+        else:
+            raise NotImplementedError(f'A {encoder} encoder is not valid.')
+        return encoder_generator
+
+    encoder_generator = get_encoder_generator(encoder)
+    y_encoder_generator = get_encoder_generator(y_encoder)
+
+    pos_encoder = args.__dict__.pop('pos_encoder')
+
+    if pos_encoder == 'none':
+        pos_encoder_generator = None
+    elif pos_encoder == 'sinus':
+        pos_encoder_generator = positional_encodings.PositionalEncoding
+    elif pos_encoder == 'learned':
+        pos_encoder_generator = positional_encodings.LearnedPositionalEncoding
+    elif pos_encoder == 'paired_scrambled_learned':
+        pos_encoder_generator = positional_encodings.PairedScrambledPositionalEncodings
+    else:
+        raise NotImplementedError(f'pos_encoer == {pos_encoder} is not valid.')
+
+    permutation_invariant_max_eval_pos = args.__dict__.pop('permutation_invariant_max_eval_pos')
+    permutation_invariant_sampling = args.__dict__.pop('permutation_invariant_sampling')
+    if permutation_invariant_max_eval_pos is not None:
+        if permutation_invariant_sampling == 'weighted':
+            get_sampler = get_weighted_single_eval_pos_sampler
+        elif permutation_invariant_sampling == 'uniform':
+            get_sampler = get_uniform_single_eval_pos_sampler
+        else:
+            raise ValueError()
+        args.__dict__['single_eval_pos_gen'] = get_sampler(permutation_invariant_max_eval_pos)
+
+
+    print("ARGS for `train`:", args.__dict__)
+
+    train(prior, criterion, encoder_generator,
+          y_encoder_generator=y_encoder_generator,pos_encoder_generator=pos_encoder_generator,
+          **args.__dict__)
+

prior-fitting/transformer.py

@@ -0,0 +1,91 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
+from torch.nn.modules.transformer import MultiheadAttention, _get_activation_fn
+
+from utils import SeqBN
+
+
+class TransformerModel(nn.Module):
+    def __init__(self, encoder, n_out, ninp, nhead, nhid, nlayers, dropout=0.0, y_encoder=None, pos_encoder=None, decoder=None, input_normalization=False):
+        super().__init__()
+        self.model_type = 'Transformer'
+        encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout, activation='gelu')
+        self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)
+        self.ninp = ninp
+        self.encoder = encoder
+        self.y_encoder = y_encoder
+        self.pos_encoder = pos_encoder
+        self.decoder = decoder(ninp, nhid, n_out) if decoder is not None else nn.Sequential(nn.Linear(ninp, nhid), nn.GELU(), nn.Linear(nhid, n_out))
+        self.input_ln = SeqBN(ninp) if input_normalization else None
+
+        self.init_weights()
+
+    @staticmethod
+    def generate_square_subsequent_mask(sz):
+        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+        return mask
+
+    @staticmethod
+    def generate_D_q_matrix(sz, query_size):
+        train_size = sz-query_size
+        mask = torch.zeros(sz,sz) == 0
+        mask[:,train_size:].zero_()
+        mask |= torch.eye(sz) == 1
+        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+        return mask
+
+    def init_weights(self):
+        initrange = 1.
+        # if isinstance(self.encoder,EmbeddingEncoder):
+        #    self.encoder.weight.data.uniform_(-initrange, initrange)
+        # self.decoder.bias.data.zero_()
+        # self.decoder.weight.data.uniform_(-initrange, initrange)
+        for layer in self.transformer_encoder.layers:
+            nn.init.zeros_(layer.linear2.weight)
+            nn.init.zeros_(layer.linear2.bias)
+            nn.init.zeros_(layer.self_attn.out_proj.weight)
+            nn.init.zeros_(layer.self_attn.out_proj.bias)
+
+    def forward(self, src, src_mask=None, single_eval_pos=None):
+        assert single_eval_pos is not None, 'Single eval pos is required now.'
+        fuse_x_y = not isinstance(src, tuple)
+        assert not(fuse_x_y and single_eval_pos is not None), \
+            'Don\'t use both fuxe_x_y and single_eval_pos (permutation equivariant setup) at the same time.'
+        if src_mask is None:
+            x_src = src if fuse_x_y else src[0]
+            if single_eval_pos is None:
+                src_mask = self.generate_square_subsequent_mask(len(x_src) if fuse_x_y else 2*len(x_src)).to(x_src.device)
+            else:
+                src_mask = self.generate_D_q_matrix(len(x_src), len(x_src)-single_eval_pos).to(x_src.device)
+        if not fuse_x_y:
+            x_src, y_src = src
+            x_src = self.encoder(x_src)
+            y_src = self.y_encoder(y_src.unsqueeze(-1))
+            if single_eval_pos is None:
+                src = torch.stack([x_src, y_src], 1).view(-1, *x_src.shape[1:])
+            else:
+                train_x = x_src[:single_eval_pos] + y_src[:single_eval_pos]
+                src = torch.cat([train_x, x_src[single_eval_pos:]], 0)
+        else:
+            src = self.encoder(src)
+
+        if self.input_ln is not None:
+            src = self.input_ln(src)
+
+        if self.pos_encoder is not None:
+            src = self.pos_encoder(src)
+
+        output = self.transformer_encoder(src, src_mask)
+        output = self.decoder(output)
+        if fuse_x_y:
+            return output
+        elif single_eval_pos is None:
+            return output[0::2]
+        else:
+            return output[single_eval_pos:]

prior-fitting/utils.py

@@ -0,0 +1,115 @@
+import math
+import argparse
+import random
+
+import torch
+from torch import nn
+from torch.optim.lr_scheduler import LambdaLR
+
+# copied from huggingface
+def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, num_cycles=0.5, last_epoch=-1):
+    """ Create a schedule with a learning rate that decreases following the
+    values of the cosine function between 0 and `pi * cycles` after a warmup
+    period during which it increases linearly between 0 and 1.
+    """
+
+    def lr_lambda(current_step):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))
+
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+# copied from huggingface
+def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
+    """
+    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after
+    a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
+
+    Args:
+        optimizer (:class:`~torch.optim.Optimizer`):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (:obj:`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (:obj:`int`):
+            The total number of training steps.
+        last_epoch (:obj:`int`, `optional`, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+
+    def lr_lambda(current_step: int):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        return max(
+            0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
+        )
+
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def get_openai_lr(transformer_model):
+    num_params = sum(p.numel() for p in transformer_model.parameters())
+    return 0.003239 - 0.0001395 * math.log(num_params)
+
+
+def get_weighted_single_eval_pos_sampler(max_len):
+    """
+    This gives a sampler that can be used for `single_eval_pos` which yields good performance for all positions p,
+    where p <= `max_len`. At most `max_len` - 1 examples are shown to the Transformer.
+    :return: Sampler that can be fed to `train()` as `single_eval_pos_gen`.
+    """
+    return lambda: random.choices(range(max_len), [1 / (max_len - i) for i in range(max_len)])[0]
+
+
+def get_uniform_single_eval_pos_sampler(max_len):
+    """
+    Just sample any evaluation position with the same weight
+    :return: Sampler that can be fed to `train()` as `single_eval_pos_gen`.
+    """
+    return lambda: random.choices(range(max_len))[0]
+
+
+class SeqBN(nn.Module):
+    def __init__(self, d_model):
+        super().__init__()
+        self.bn = nn.BatchNorm1d(d_model)
+        self.d_model = d_model
+
+    def forward(self, x):
+        assert self.d_model == x.shape[-1]
+        flat_x = x.view(-1, self.d_model)
+        flat_x = self.bn(flat_x)
+        return flat_x.view(*x.shape)
+
+
+def set_locals_in_self(locals):
+    self = locals['self']
+    for var_name, val in locals.items():
+        if var_name != 'self': setattr(self, var_name, val)
+
+
+default_device = 'cuda:0' if torch.cuda.is_available() else 'cpu:0'
+
+
+# Copied from StackOverflow, but we do an eval on the values additionally
+class StoreDictKeyPair(argparse.Action):
+    def __init__(self, option_strings, dest, nargs=None, **kwargs):
+        self._nargs = nargs
+        super(StoreDictKeyPair, self).__init__(option_strings, dest, nargs=nargs, **kwargs)
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        my_dict = {}
+        for kv in values:
+            k, v = kv.split("=")
+            try:
+                my_dict[k] = eval(v)
+            except NameError:
+                my_dict[k] = v
+        setattr(namespace, self.dest, my_dict)
+        print("dict values: {}".format(my_dict))
+
+

requirements.txt

@@ -0,0 +1,13 @@
+# Recommend to use python >= 3.9
+gpytorch==1.5.0
+pyro-ppl==1.7.0
+torch==1.9.0
+scikit-learn==0.24.2
+pyyaml==5.4.1
+blitz-bayesian-pytorch==0.2.7
+seaborn==0.11.2
+xgboost==1.4.0
+tqdm==4.62.1
+numpy==1.21.2
+openml==0.12.2
+catboost==0.26.1