```python
import torch
import torch.nn.functional as F

import torchsde
import math
import matplotlib.pyplot as plt

import numpy as np

from tqdm.notebook import tqdm

from torch import _vmap_internals
```

```python
cd ..
```

    C:\Users\vargf\OneDrive\Documents\Projects\ControlledFollMerDrift

```python
from cfollmer.objectives import log_g, relative_entropy_control_cost, stl_relative_entropy_control_cost
from cfollmer.sampler_utils import FollmerSDE
from cfollmer.trainers import basic_batched_trainer
```

# The Model

\begin{align}
\theta &\sim \mathcal{N}(\theta | 0, \sigma_w^2 \mathbb{I}) \\
y_i | x_i, \theta &\sim  \mathrm{Bernouli}\left[\mathrm{NN}_{\theta}\left(x_i \right)\right]
\end{align}

We want samples from $p(\theta | \{(y_i, x_i)\})$. Note $f(x; \theta)$ is a neural net with params $\theta$

## Loading the iris dataset

```python
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

device = "cuda" if torch.cuda.is_available() else "cpu"

iris = load_iris()
X = iris['data']
y = iris['target']

# Binary classification
X = X[~(y==2)][:,[0,1]]
y = y[~(y==2)]

# dummy dims
X = np.concatenate((torch.ones(X.shape[0],1), torch.tensor(X) ), axis=1)

names = iris['target_names']
feature_names = iris['feature_names']

# Scale data to have mean 0 and variance 1
# which is importance for convergence of the neural network
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# Split the data set into training and testing
X_train, X_test, y_train, y_test = train_test_split(
    X_scaled, y, test_size=0.2, random_state=2)

X_train, X_test, y_train, y_test = \
    torch.tensor(X_train, dtype=torch.float32, device=device), \
    torch.tensor(X_test, dtype=torch.float32, device=device), \
    torch.tensor(y_train, dtype=torch.float32, device=device), \
    torch.tensor(y_test, dtype=torch.float32, device=device)
```

```python
fig, ax1 = plt.subplots(1, 1, figsize=(16, 6))
for target, target_name in enumerate(names[0:2]):
    X_plot = X[y == target]
    ax1.plot(X_plot[:, 1], X_plot[:, 2],
             linestyle='none',
             marker='o',
             label=target_name)
ax1.set_xlabel(feature_names[0])
ax1.set_ylabel(feature_names[1])
ax1.axis('equal')
ax1.legend();
```

$$\DeclareMathOperator*{\argmin}{arg\,min}$$
$$\def\E{{\mathbb{E}}}$$
$$\def\rvu{{\mathbf{u}}}$$
$$\def\rvTheta{{\bm{\Theta}}}$$
$$\def\gU{{\mathcal{U}}}$$
$$\def\mX{{\mathbf{X}}}$$

## Controlled Schrodinger Follmer Sampler

The objevtive we are trying to implement is:

\begin{align}
  \mathbf{u}_t^{*}=  \argmin_{\rvu_t \in \mathcal{U}}\mathbb{E}\left[\frac{1}{2\gamma}\int_0^1||\rvu(t, \Theta_t)||^2 dt - \ln\left(\frac{ p(\mX | \Theta_1)p(\Theta_1)}{\mathcal{N}(\Theta_1|\mathbf{0}, \gamma \mathbb{I} )}\right)\right] \
\end{align}

Where:
\begin{align}
d\Theta_t = \rvu(t, \Theta_t)dt + \sqrt{\gamma} dB_t
\end{align}

To do so we use the EM discretisation.

```python
import torch.nn.functional as F

class OnedRegressionForwardNet(object):

    def __init__(
        self, input_dim=1, output_dim=1, depth=None,
        width=20, width_seq=None, device="cpu", activation=F.relu
    ):

        self.device = device
        self.output_dim = output_dim
        self.input_dim = input_dim
        self.activation = activation

        self.depth = depth
        if not self.depth:
            self.depth = 1
        if not width_seq:
            self.width = width
            self.width_seq = [self.width] * (self.depth + 1)
            self.shapes = [(self.width_seq[i-1], self.width_seq[i])  for i in range(1,self.depth)]
            self.shapes += [(self.width_seq[-1], self.output_dim)]
            self.shapes = [(self.input_dim, self.width_seq[0])] + self.shapes

        self.dim = sum([wx * wy + wy for wx, wy in self.shapes])

    def forward(self, x, Θ):
        index = 0
        n, d = x.shape

#         dim_bl =  sum([wx * wy + wy for wx, wy in self.shapes[:-1]])
#         Θ[:dim_bl] = (Θ[:dim_bl] - Θ[:dim_bl].mean()) / Θ[:dim_bl].std()
#         σ_Θ, μ_Θ = Θ.std(), Θ.mean()
#         Θ = (Θ - μ_Θ) / σ_Θ

        for wx, wy in self.shapes[:-1]:
            x = F.linear(
                x,
                Θ[index: index + wx * wy].reshape(wy, wx),
                Θ[index + wx * wy: index + wx * wy + wy].reshape(1,wy)
            )
            x = self.activation(x)
            index += wx * wy  + wy
        wx, wy = self.shapes[-1]
        x = F.linear(
            x,
            Θ[index: index + wx * wy].reshape(wy, wx), #* σ_Θ + μ_Θ,
            Θ[index + wx * wy: index + wx * wy + wy].reshape(1,wy) # * σ_Θ + μ_Θ
        )
        return x.to(self.device)

    def map_forward(self, x, Θ):
        preds_func = lambda θ: self.forward(x, θ)
        batched_preds = torch._vmap_internals.vmap(preds_func)
        preds = torch.hstack(list(map(preds_func, Θ)))
        return preds

```

```python
net = OnedRegressionForwardNet(
    3,1, device=device, depth=1, width=30, activation=F.tanh
)

def gaussian_prior(Θ, σ_w=1.8):
    """
    Logistic regresion bayesian prior
    """
    return -0.5 * (Θ**2).sum(axis=1) / σ_w

# def log_likelihood_vmap(Θ, X, y):
#     """
#     Hoping this implementation is less buggy / faster

#     still feels a bit slow.
#     """
#     logits = X.mm(Θ.T)

#     pos_weights = torch.ones(logits.shape[0], device=device)

#     loss = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weights, reduction="sum")

#     # TODO: Double check this is right, changed to a minus sign here
#     loss_ = lambda x: -1.0 * loss(x, y)

#     batched_loss =  torch._vmap_internals.vmap(loss_)

#     return batched_loss(logits.T)

def log_likelihood_vmap_nn(Θ, X, y, net=net):
    """
    Hoping this implementation is less buggy / faster

    still feels a bit slow.
    """
    pos_weights = torch.ones(X.shape[0], device=device)

    def loss(θ):
        preds = net.forward(X, θ)
        bce = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weights, reduction="sum")
        ll_bcs = -1.0 * bce(preds.reshape(-1), y.reshape(-1))
        return ll_bcs

    batched_loss =  torch._vmap_internals.vmap(loss)

    return batched_loss(Θ)
```

```python
γ = 0.2
Δt=0.05

dim= net.dim

sde, losses = basic_batched_trainer(
    γ, Δt, gaussian_prior, log_likelihood_vmap_nn, dim, X_train, y_train,
    method="euler", stl=True, adjoint=False, optimizer=None,
    num_steps=400, batch_size_data=X_train.shape[0], batch_size_Θ=20,
    batchnorm=True, device=device, lr=0.001
)
```

    LOS STL?: <function stl_relative_entropy_control_cost at 0x0000023887981CA0>

      0%|          | 0/400 [00:00<?, ?it/s]

      0%|          | 0/1 [00:00<?, ?it/s]

    0.6142825484275818

    C:\Users\vargf\OneDrive\Documents\Projects\ControlledFollMerDrift\cfollmer\objectives.py:135: UserWarning: torch.vmap is an experimental prototype that is subject to change and/or deletion. Please use at your own risk. There may be unexpected performance cliffs due to certain operators not being implemented. To see detailed performance warnings please use `torch._C._debug_only_display_vmap_fallback_warnings(True) before the call to `vmap`.
      f = _vmap_internals.vmap(f_)
    C:\Users\vargf\OneDrive\Documents\Projects\ControlledFollMerDrift\cfollmer\objectives.py:136: UserWarning: torch.vmap is an experimental prototype that is subject to change and/or deletion. Please use at your own risk. There may be unexpected performance cliffs due to certain operators not being implemented. To see detailed performance warnings please use `torch._C._debug_only_display_vmap_fallback_warnings(True) before the call to `vmap`.
      f_detached = _vmap_internals.vmap(sde.f_detached)
    C:\Users\vargf\AppData\Local\Temp/ipykernel_26468/579694690.py:47: UserWarning: torch.vmap is an experimental prototype that is subject to change and/or deletion. Please use at your own risk. There may be unexpected performance cliffs due to certain operators not being implemented. To see detailed performance warnings please use `torch._C._debug_only_display_vmap_fallback_warnings(True) before the call to `vmap`.
      batched_loss =  torch._vmap_internals.vmap(loss)
    C:\Users\vargf\anaconda3\lib\site-packages\torch\nn\functional.py:1795: UserWarning: nn.functional.tanh is deprecated. Use torch.tanh instead.
      warnings.warn("nn.functional.tanh is deprecated. Use torch.tanh instead.")

      0%|          | 0/1 [00:00<?, ?it/s]

    0.5484275817871094

      0%|          | 0/1 [00:00<?, ?it/s]

    0.4473325312137604

      0%|          | 0/1 [00:00<?, ?it/s]

    0.4502784311771393

      0%|          | 0/1 [00:00<?, ?it/s]

    0.36465150117874146

      0%|          | 0/1 [00:00<?, ?it/s]

    0.3416145145893097

      0%|          | 0/1 [00:00<?, ?it/s]

    0.3479955494403839

      0%|          | 0/1 [00:00<?, ?it/s]

    0.2824741303920746

      0%|          | 0/1 [00:00<?, ?it/s]

    0.21983222663402557

      0%|          | 0/1 [00:00<?, ?it/s]

    0.20826180279254913

      0%|          | 0/1 [00:00<?, ?it/s]

    0.21422818303108215

      0%|          | 0/1 [00:00<?, ?it/s]

    0.1416165679693222

      0%|          | 0/1 [00:00<?, ?it/s]

    0.1280628889799118

      0%|          | 0/1 [00:00<?, ?it/s]

    0.05953091010451317

      0%|          | 0/1 [00:00<?, ?it/s]

    0.08050859719514847

      0%|          | 0/1 [00:00<?, ?it/s]

    0.08609960973262787

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.006078376434743404

      0%|          | 0/1 [00:00<?, ?it/s]

    0.009355875663459301

      0%|          | 0/1 [00:00<?, ?it/s]

    0.016610614955425262

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.036033034324645996

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.046597935259342194

      0%|          | 0/1 [00:00<?, ?it/s]

    0.008752009831368923

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03596426546573639

      0%|          | 0/1 [00:00<?, ?it/s]

    0.013041503727436066

      0%|          | 0/1 [00:00<?, ?it/s]

    0.007062550634145737

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.14388732612133026

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07516121864318848

      0%|          | 0/1 [00:00<?, ?it/s]

    0.011232366785407066

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06389092653989792

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.10368144512176514

      0%|          | 0/1 [00:00<?, ?it/s]

    0.025115717202425003

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05433184653520584

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0793132483959198

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.15328392386436462

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03840652108192444

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06193932518362999

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.031466808170080185

      0%|          | 0/1 [00:00<?, ?it/s]

    0.005405769683420658

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.11160796135663986

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0876423567533493

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0436294861137867

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.11185093224048615

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0303993858397007

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0023379623889923096

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.039690837264060974

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.01715460978448391

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05997025594115257

      0%|          | 0/1 [00:00<?, ?it/s]

    0.029942471534013748

      0%|          | 0/1 [00:00<?, ?it/s]

    0.01716366969048977

      0%|          | 0/1 [00:00<?, ?it/s]

    0.007308214902877808

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.11066372692584991

      0%|          | 0/1 [00:00<?, ?it/s]

    0.06905031949281693

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.08778326213359833

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06758280843496323

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.013384580612182617

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.09507667273283005

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.016828328371047974

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03416025638580322

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07367303222417831

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06724090874195099

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.028630254790186882

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.017252467572689056

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03203723207116127

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.011862092651426792

      0%|          | 0/1 [00:00<?, ?it/s]

    0.048885613679885864

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.005813518073409796

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.01020615641027689

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0847427025437355

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.01338405441492796

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.09887602180242538

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.053541410714387894

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.011121721938252449

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.042205821722745895

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0074891201220452785

      0%|          | 0/1 [00:00<?, ?it/s]

    0.03184041008353233

      0%|          | 0/1 [00:00<?, ?it/s]

    0.024172024801373482

      0%|          | 0/1 [00:00<?, ?it/s]

    0.007109713740646839

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07229097932577133

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05140851065516472

      0%|          | 0/1 [00:00<?, ?it/s]

    0.01891208626329899

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0019971539732068777

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.012247535400092602

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0756739154458046

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05858224630355835

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0680375024676323

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07439880818128586

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.007914464920759201

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07552327960729599

      0%|          | 0/1 [00:00<?, ?it/s]

    0.014828534796833992

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06831517815589905

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.036100830882787704

      0%|          | 0/1 [00:00<?, ?it/s]

    0.03302161023020744

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03397815302014351

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06191062927246094

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.011270126327872276

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07995075732469559

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.042087897658348083

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.036574479192495346

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04629189148545265

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.02194184623658657

      0%|          | 0/1 [00:00<?, ?it/s]

    0.0789874717593193

      0%|          | 0/1 [00:00<?, ?it/s]

    0.02054011821746826

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.013228652067482471

      0%|          | 0/1 [00:00<?, ?it/s]

    0.05409296229481697

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.02798142470419407

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04969160258769989

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.057179730385541916

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06178957223892212

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.020434420555830002

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07913928478956223

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.08944673091173172

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05906078219413757

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.009511872194707394

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05294156074523926

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0008628565119579434

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.10693657398223877

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.020187102258205414

      0%|          | 0/1 [00:00<?, ?it/s]

    0.0073417299427092075

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.10207366198301315

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07574330270290375

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.015799425542354584

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.008383825421333313

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.08340520411729813

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.008825652301311493

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07460209727287292

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06280653923749924

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.13905757665634155

      0%|          | 0/1 [00:00<?, ?it/s]

    0.00032576979720033705

      0%|          | 0/1 [00:00<?, ?it/s]

    0.03859635815024376

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.036444179713726044

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04360423982143402

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03866339847445488

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05000493675470352

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03882714360952377

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.02973007969558239

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0330064557492733

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.125828355550766

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.006134114693850279

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.12021093815565109

      0%|          | 0/1 [00:00<?, ?it/s]

    0.009758037514984608

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0055194213055074215

      0%|          | 0/1 [00:00<?, ?it/s]

    0.022512119263410568

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03906669840216637

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07385485619306564

      0%|          | 0/1 [00:00<?, ?it/s]

    0.044757992029190063

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04435334727168083

      0%|          | 0/1 [00:00<?, ?it/s]

    0.0022190362215042114

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.08005756884813309

      0%|          | 0/1 [00:00<?, ?it/s]

    0.007960936054587364

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04746043309569359

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07163546979427338

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.023875070735812187

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.027583498507738113

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.037111781537532806

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.053783219307661057

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0685693547129631

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.02491629496216774

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06057020649313927

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05690523609519005

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.029619647189974785

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.012864597141742706

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.030755018815398216

      0%|          | 0/1 [00:00<?, ?it/s]

    0.030750632286071777

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05165916681289673

      0%|          | 0/1 [00:00<?, ?it/s]

    0.025617359206080437

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05048106983304024

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.017077261582016945

      0%|          | 0/1 [00:00<?, ?it/s]

    0.0018924266332760453

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.054924167692661285

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.052407026290893555

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.005650452803820372

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04258588328957558

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0759638324379921

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05401661992073059

      0%|          | 0/1 [00:00<?, ?it/s]

    0.018218087032437325

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.01713041216135025

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06998025625944138

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04257845878601074

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.15254752337932587

      0%|          | 0/1 [00:00<?, ?it/s]

    0.004868763964623213

      0%|          | 0/1 [00:00<?, ?it/s]

    0.02679472044110298

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05136943608522415

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.025690734386444092

      0%|          | 0/1 [00:00<?, ?it/s]

    0.023264169692993164

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.056069452315568924

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03386129438877106

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07515844702720642

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0671985074877739

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.003848437685519457

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07138577848672867

      0%|          | 0/1 [00:00<?, ?it/s]

    0.0363314189016819

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03786192089319229

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06379406899213791

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03534262999892235

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07531312853097916

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04345773905515671

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.061422526836395264

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04645416513085365

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04911838471889496

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07713698595762253

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.011969472281634808

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.08831439167261124

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.016369013115763664

      0%|          | 0/1 [00:00<?, ?it/s]

    0.0017520649125799537

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05089762806892395

      0%|          | 0/1 [00:00<?, ?it/s]

    0.02125023491680622

      0%|          | 0/1 [00:00<?, ?it/s]

    0.06981562823057175

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.012566273100674152

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.026926100254058838

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.005359499715268612

      0%|          | 0/1 [00:00<?, ?it/s]

    0.012591840699315071

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06857622414827347

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03760841488838196

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.011333378963172436

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.010964902117848396

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.057309336960315704

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05350165441632271

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04743897542357445

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0043594432063400745

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.053854890167713165

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0062580169178545475

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04426759108901024

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05407492071390152

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.024996887892484665

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06389033049345016

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.08963315188884735

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0524543933570385

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.12244751304388046

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.00635957345366478

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0358552448451519

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.027749160304665565

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.09038262069225311

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06435742229223251

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0027135456912219524

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05032682418823242

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07365413755178452

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06998325884342194

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0778547152876854

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06616618484258652

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05663363263010979

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.019634483382105827

      0%|          | 0/1 [00:00<?, ?it/s]

    0.04151206091046333

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06716714799404144

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03538958355784416

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.08218108862638474

      0%|          | 0/1 [00:00<?, ?it/s]

    0.004701628815382719

      0%|          | 0/1 [00:00<?, ?it/s]

    0.0017469382146373391

      0%|          | 0/1 [00:00<?, ?it/s]

    0.01499372161924839

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.010968988761305809

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07785826921463013

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.058221783488988876

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04647994786500931

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.044243279844522476

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0652492493391037

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.058973606675863266

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.10066777467727661

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05275845527648926

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0569780170917511

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.034446991980075836

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.09686260670423508

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.025469660758972168

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07912570983171463

      0%|          | 0/1 [00:00<?, ?it/s]

    0.009209475480020046

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.006181337404996157

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.08965963870286942

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03345806524157524

      0%|          | 0/1 [00:00<?, ?it/s]

    0.051917996257543564

      0%|          | 0/1 [00:00<?, ?it/s]

    0.004713513422757387

      0%|          | 0/1 [00:00<?, ?it/s]

    0.05096372961997986

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0652390867471695

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.013441065326333046

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.029762575402855873

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04977777972817421

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.044698651880025864

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.024665068835020065

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04512549564242363

      0%|          | 0/1 [00:00<?, ?it/s]

    0.04135502129793167

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04149525985121727

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0031684041023254395

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.025288019329309464

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07295200228691101

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.034475285559892654

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0880858525633812

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0781109482049942

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.10871180146932602

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.13549646735191345

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06240592151880264

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07178980857133865

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04729149863123894

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.026618918403983116

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.02523321844637394

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.09059762954711914

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04067535325884819

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03673673793673515

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04453755542635918

      0%|          | 0/1 [00:00<?, ?it/s]

    0.02002827264368534

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.00475192628800869

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06165891885757446

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.009200246073305607

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0030528707429766655

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05331080034375191

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.1422048956155777

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06716674566268921

      0%|          | 0/1 [00:00<?, ?it/s]

    0.006956851575523615

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06973776966333389

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0654275119304657

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06472715735435486

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0334104523062706

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.029803713783621788

      0%|          | 0/1 [00:00<?, ?it/s]

    0.015631506219506264

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.08096999675035477

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06623058021068573

      0%|          | 0/1 [00:00<?, ?it/s]

    0.0367695651948452

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.08316795527935028

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.012080232612788677

      0%|          | 0/1 [00:00<?, ?it/s]

    0.013144010677933693

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.044727873057127

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03981294482946396

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.008867970667779446

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.007173065096139908

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06874512135982513

      0%|          | 0/1 [00:00<?, ?it/s]

    0.02353695221245289

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07540923357009888

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.026348648592829704

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.023197071626782417

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.11383236944675446

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06648343056440353

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07975773513317108

      0%|          | 0/1 [00:00<?, ?it/s]

    0.003525529056787491

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03923170641064644

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.028089020401239395

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.08751492202281952

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07536058872938156

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04172651097178459

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.026523906737565994

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03351398929953575

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.03464892879128456

      0%|          | 0/1 [00:00<?, ?it/s]

    0.03887905552983284

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05492640659213066

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.006427271757274866

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.062457866966724396

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.036685761064291

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07170569151639938

      0%|          | 0/1 [00:00<?, ?it/s]

    0.01574891060590744

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.011154123581945896

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.08681545406579971

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.018573932349681854

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05427476391196251

      0%|          | 0/1 [00:00<?, ?it/s]

    0.02066950313746929

      0%|          | 0/1 [00:00<?, ?it/s]

    0.007995781488716602

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07969412207603455

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.061624784022569656

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07458213716745377

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.01418994553387165

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.02757524512708187

      0%|          | 0/1 [00:00<?, ?it/s]

    0.004010220989584923

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.020229294896125793

      0%|          | 0/1 [00:00<?, ?it/s]

    0.028411591425538063

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.013759388588368893

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.046712055802345276

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.10412490367889404

      0%|          | 0/1 [00:00<?, ?it/s]

    0.03685170039534569

      0%|          | 0/1 [00:00<?, ?it/s]

    0.0014291101833805442

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06345359981060028

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06264366954565048

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04861581325531006

      0%|          | 0/1 [00:00<?, ?it/s]

    0.02435915544629097

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.01793580688536167

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0637972354888916

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.014055440202355385

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04714873805642128

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0065234326757490635

      0%|          | 0/1 [00:00<?, ?it/s]

    0.07194950431585312

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07455312460660934

      0%|          | 0/1 [00:00<?, ?it/s]

    0.008192349225282669

      0%|          | 0/1 [00:00<?, ?it/s]

    0.029232880100607872

      0%|          | 0/1 [00:00<?, ?it/s]

    0.003191909985616803

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.014510511420667171

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.005821959115564823

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.09569769352674484

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.01241032499819994

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.1129368469119072

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.0415160171687603

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.025916293263435364

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.05111085996031761

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06516343355178833

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.025754520669579506

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.044067032635211945

      0%|          | 0/1 [00:00<?, ?it/s]

    0.006671845447272062

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.06856798380613327

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04689745232462883

      0%|          | 0/1 [00:00<?, ?it/s]

    0.008680857717990875

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.07118958979845047

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.02640773169696331

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.10959835350513458

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.04767514392733574

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.051034945994615555

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.01595236174762249

      0%|          | 0/1 [00:00<?, ?it/s]

    0.000419425981817767

      0%|          | 0/1 [00:00<?, ?it/s]

    -0.047449443489313126

```python
losses
```

    [tensor(0.6143),
     tensor(0.5484),
     tensor(0.4473),
     tensor(0.4503),
     tensor(0.3647),
     tensor(0.3416),
     tensor(0.3480),
     tensor(0.2825),
     tensor(0.2198),
     tensor(0.2083),
     tensor(0.2142),
     tensor(0.1416),
     tensor(0.1281),
     tensor(0.0595),
     tensor(0.0805),
     tensor(0.0861),
     tensor(-0.0061),
     tensor(0.0094),
     tensor(0.0166),
     tensor(-0.0360),
     tensor(-0.0466),
     tensor(0.0088),
     tensor(-0.0360),
     tensor(0.0130),
     tensor(0.0071),
     tensor(-0.1439),
     tensor(-0.0752),
     tensor(0.0112),
     tensor(-0.0639),
     tensor(-0.1037),
     tensor(0.0251),
     tensor(-0.0543),
     tensor(-0.0793),
     tensor(-0.1533),
     tensor(-0.0384),
     tensor(-0.0619),
     tensor(-0.0315),
     tensor(0.0054),
     tensor(-0.1116),
     tensor(-0.0876),
     tensor(-0.0436),
     tensor(-0.1119),
     tensor(-0.0304),
     tensor(-0.0023),
     tensor(-0.0397),
     tensor(-0.0172),
     tensor(-0.0600),
     tensor(0.0299),
     tensor(0.0172),
     tensor(0.0073),
     tensor(-0.1107),
     tensor(0.0691),
     tensor(-0.0878),
     tensor(-0.0676),
     tensor(-0.0134),
     tensor(-0.0951),
     tensor(-0.0168),
     tensor(-0.0342),
     tensor(-0.0737),
     tensor(-0.0672),
     tensor(-0.0286),
     tensor(-0.0173),
     tensor(-0.0320),
     tensor(-0.0119),
     tensor(0.0489),
     tensor(-0.0058),
     tensor(-0.0102),
     tensor(-0.0847),
     tensor(-0.0134),
     tensor(-0.0989),
     tensor(-0.0535),
     tensor(-0.0111),
     tensor(-0.0422),
     tensor(-0.0075),
     tensor(0.0318),
     tensor(0.0242),
     tensor(0.0071),
     tensor(-0.0723),
     tensor(-0.0514),
     tensor(0.0189),
     tensor(-0.0020),
     tensor(-0.0122),
     tensor(-0.0757),
     tensor(-0.0586),
     tensor(-0.0680),
     tensor(-0.0744),
     tensor(-0.0079),
     tensor(-0.0755),
     tensor(0.0148),
     tensor(-0.0683),
     tensor(-0.0361),
     tensor(0.0330),
     tensor(-0.0340),
     tensor(-0.0619),
     tensor(-0.0113),
     tensor(-0.0800),
     tensor(-0.0421),
     tensor(-0.0366),
     tensor(-0.0463),
     tensor(-0.0219),
     tensor(0.0790),
     tensor(0.0205),
     tensor(-0.0132),
     tensor(0.0541),
     tensor(-0.0280),
     tensor(-0.0497),
     tensor(-0.0572),
     tensor(-0.0618),
     tensor(-0.0204),
     tensor(-0.0791),
     tensor(-0.0894),
     tensor(-0.0591),
     tensor(-0.0095),
     tensor(-0.0529),
     tensor(-0.0009),
     tensor(-0.1069),
     tensor(-0.0202),
     tensor(0.0073),
     tensor(-0.1021),
     tensor(-0.0757),
     tensor(-0.0158),
     tensor(-0.0084),
     tensor(-0.0834),
     tensor(-0.0088),
     tensor(-0.0746),
     tensor(-0.0628),
     tensor(-0.1391),
     tensor(0.0003),
     tensor(0.0386),
     tensor(-0.0364),
     tensor(-0.0436),
     tensor(-0.0387),
     tensor(-0.0500),
     tensor(-0.0388),
     tensor(-0.0297),
     tensor(-0.0330),
     tensor(-0.1258),
     tensor(-0.0061),
     tensor(-0.1202),
     tensor(0.0098),
     tensor(-0.0055),
     tensor(0.0225),
     tensor(-0.0391),
     tensor(-0.0739),
     tensor(0.0448),
     tensor(-0.0444),
     tensor(0.0022),
     tensor(-0.0801),
     tensor(0.0080),
     tensor(-0.0475),
     tensor(-0.0716),
     tensor(-0.0239),
     tensor(-0.0276),
     tensor(-0.0371),
     tensor(-0.0538),
     tensor(-0.0686),
     tensor(-0.0249),
     tensor(-0.0606),
     tensor(-0.0569),
     tensor(-0.0296),
     tensor(-0.0129),
     tensor(-0.0308),
     tensor(0.0308),
     tensor(-0.0517),
     tensor(0.0256),
     tensor(-0.0505),
     tensor(-0.0171),
     tensor(0.0019),
     tensor(-0.0549),
     tensor(-0.0524),
     tensor(-0.0057),
     tensor(-0.0426),
     tensor(-0.0760),
     tensor(-0.0540),
     tensor(0.0182),
     tensor(-0.0171),
     tensor(-0.0700),
     tensor(-0.0426),
     tensor(-0.1525),
     tensor(0.0049),
     tensor(0.0268),
     tensor(-0.0514),
     tensor(-0.0257),
     tensor(0.0233),
     tensor(-0.0561),
     tensor(-0.0339),
     tensor(-0.0752),
     tensor(-0.0672),
     tensor(-0.0038),
     tensor(-0.0714),
     tensor(0.0363),
     tensor(-0.0379),
     tensor(-0.0638),
     tensor(-0.0353),
     tensor(-0.0753),
     tensor(-0.0435),
     tensor(-0.0614),
     tensor(-0.0465),
     tensor(-0.0491),
     tensor(-0.0771),
     tensor(-0.0120),
     tensor(-0.0883),
     tensor(-0.0164),
     tensor(0.0018),
     tensor(-0.0509),
     tensor(0.0213),
     tensor(0.0698),
     tensor(-0.0126),
     tensor(-0.0269),
     tensor(-0.0054),
     tensor(0.0126),
     tensor(-0.0686),
     tensor(-0.0376),
     tensor(-0.0113),
     tensor(-0.0110),
     tensor(-0.0573),
     tensor(-0.0535),
     tensor(-0.0474),
     tensor(-0.0044),
     tensor(-0.0539),
     tensor(-0.0063),
     tensor(-0.0443),
     tensor(-0.0541),
     tensor(-0.0250),
     tensor(-0.0639),
     tensor(-0.0896),
     tensor(-0.0525),
     tensor(-0.1224),
     tensor(-0.0064),
     tensor(-0.0359),
     tensor(-0.0277),
     tensor(-0.0904),
     tensor(-0.0644),
     tensor(-0.0027),
     tensor(-0.0503),
     tensor(-0.0737),
     tensor(-0.0700),
     tensor(-0.0779),
     tensor(-0.0662),
     tensor(-0.0566),
     tensor(-0.0196),
     tensor(0.0415),
     tensor(-0.0672),
     tensor(-0.0354),
     tensor(-0.0822),
     tensor(0.0047),
     tensor(0.0017),
     tensor(0.0150),
     tensor(-0.0110),
     tensor(-0.0779),
     tensor(-0.0582),
     tensor(-0.0465),
     tensor(-0.0442),
     tensor(-0.0652),
     tensor(-0.0590),
     tensor(-0.1007),
     tensor(-0.0528),
     tensor(-0.0570),
     tensor(-0.0344),
     tensor(-0.0969),
     tensor(-0.0255),
     tensor(-0.0791),
     tensor(0.0092),
     tensor(-0.0062),
     tensor(-0.0897),
     tensor(-0.0335),
     tensor(0.0519),
     tensor(0.0047),
     tensor(0.0510),
     tensor(-0.0652),
     tensor(-0.0134),
     tensor(-0.0298),
     tensor(-0.0498),
     tensor(-0.0447),
     tensor(-0.0247),
     tensor(-0.0451),
     tensor(0.0414),
     tensor(-0.0415),
     tensor(-0.0032),
     tensor(-0.0253),
     tensor(-0.0730),
     tensor(-0.0345),
     tensor(-0.0881),
     tensor(-0.0781),
     tensor(-0.1087),
     tensor(-0.1355),
     tensor(-0.0624),
     tensor(-0.0718),
     tensor(-0.0473),
     tensor(-0.0266),
     tensor(-0.0252),
     tensor(-0.0906),
     tensor(-0.0407),
     tensor(-0.0367),
     tensor(-0.0445),
     tensor(0.0200),
     tensor(-0.0048),
     tensor(-0.0617),
     tensor(-0.0092),
     tensor(-0.0031),
     tensor(-0.0533),
     tensor(-0.1422),
     tensor(-0.0672),
     tensor(0.0070),
     tensor(-0.0697),
     tensor(-0.0654),
     tensor(-0.0647),
     tensor(-0.0334),
     tensor(-0.0298),
     tensor(0.0156),
     tensor(-0.0810),
     tensor(-0.0662),
     tensor(0.0368),
     tensor(-0.0832),
     tensor(-0.0121),
     tensor(0.0131),
     tensor(-0.0447),
     tensor(-0.0398),
     tensor(-0.0089),
     tensor(-0.0072),
     tensor(-0.0687),
     tensor(0.0235),
     tensor(-0.0754),
     tensor(-0.0263),
     tensor(-0.0232),
     tensor(-0.1138),
     tensor(-0.0665),
     tensor(-0.0798),
     tensor(0.0035),
     tensor(-0.0392),
     tensor(-0.0281),
     tensor(-0.0875),
     tensor(-0.0754),
     tensor(-0.0417),
     tensor(-0.0265),
     tensor(-0.0335),
     tensor(-0.0346),
     tensor(0.0389),
     tensor(-0.0549),
     tensor(-0.0064),
     tensor(-0.0625),
     tensor(-0.0367),
     tensor(-0.0717),
     tensor(0.0157),
     tensor(-0.0112),
     tensor(-0.0868),
     tensor(-0.0186),
     tensor(-0.0543),
     tensor(0.0207),
     tensor(0.0080),
     tensor(-0.0797),
     tensor(-0.0616),
     tensor(-0.0746),
     tensor(-0.0142),
     tensor(-0.0276),
     tensor(0.0040),
     tensor(-0.0202),
     tensor(0.0284),
     tensor(-0.0138),
     tensor(-0.0467),
     tensor(-0.1041),
     tensor(0.0369),
     tensor(0.0014),
     tensor(-0.0635),
     tensor(-0.0626),
     tensor(-0.0486),
     tensor(0.0244),
     tensor(-0.0179),
     tensor(-0.0638),
     tensor(-0.0141),
     tensor(-0.0471),
     tensor(-0.0065),
     tensor(0.0719),
     tensor(-0.0746),
     tensor(0.0082),
     tensor(0.0292),
     tensor(0.0032),
     tensor(-0.0145),
     tensor(-0.0058),
     tensor(-0.0957),
     tensor(-0.0124),
     tensor(-0.1129),
     tensor(-0.0415),
     tensor(-0.0259),
     tensor(-0.0511),
     tensor(-0.0652),
     tensor(-0.0258),
     tensor(-0.0441),
     tensor(0.0067),
     tensor(-0.0686),
     tensor(-0.0469),
     tensor(0.0087),
     tensor(-0.0712),
     tensor(-0.0264),
     tensor(-0.1096),
     tensor(-0.0477),
     tensor(-0.0510),
     tensor(-0.0160),
     tensor(0.0004),
     tensor(-0.0474)]

```python
plt.plot(losses[:])
```

```python
X_train.shape
```

    torch.Size([80, 3])

```python
t_size = int(math.ceil(1.0/Δt))
ts = torch.linspace(0, 1, t_size).to(device)
no_posterior_samples = 100
Θ_0 = torch.zeros((no_posterior_samples, net.dim)).to(device)

Θ_1 = torchsde.sdeint(sde, Θ_0, ts, dt=Δt)[-1,...]
```

```python
fig, (ax1,ax2,ax3) = plt.subplots(1,3)

ax1.hist(Θ_1[:,0].cpu().detach().numpy())
ax2.hist(Θ_1[:,1].cpu().detach().numpy())
ax3.hist(Θ_1[:,2].cpu().detach().numpy())
```

```python
def predc(X, Θ):
    return torch.vstack([torch.sigmoid(net.forward(X, θ)[None,...]) for θ in Θ]).mean(dim=0)
```

```python
pred = predc(X_train, Θ_1)
```

```python

((pred > 0.5).float().flatten()== y_train).float().mean()
```

    tensor(0.9875, device='cuda:0')

```python
pred_test = predc(X_test.float(), Θ_1)
```

```python
((pred_test > 0.5).float().flatten() == y_test).float().mean()
```

    tensor(1., device='cuda:0')

```python
y_test, (pred_test < 0.5).long().flatten()
```

    (tensor([1., 0., 1., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 0., 0., 0., 0., 0.,
             1., 1.], device='cuda:0'),
     tensor([0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0],
            device='cuda:0'))

```python
plt.clf()

scaler = StandardScaler()
X_scaled2 = scaler.fit_transform(X)

plt.plot(X_scaled2[y==1, 1], X_scaled2[y==1, 2], 'bx')
plt.plot(X_scaled2[y==0, 1], X_scaled2[y==0, 2], 'ro')
plt.legend(('y=1', 'y=0'))

# # Overlay contour plot of approximate predictive distribution:
x_grid = np.arange(-4, 4, 0.009)
X1, X2 = np.meshgrid(x_grid, x_grid)
NG = X1.size
X_test_2 = np.hstack((np.ones((NG,1)), X1.reshape(NG,1), X2.reshape(NG,1)))
X_test_2.shape

X_test_2_tt = torch.tensor(X_test_2).to(device).float()
p_test  = (predc(X_test_2_tt.cpu(),Θ_1.cpu())).mean(axis=1).detach().cpu().numpy()

# kappa = 1.0 / np.sqrt(1 + (np.pi/8)*np.sum(np.dot(X_test,V)*X_test, 1))
# p_test = 1.0 / (1+np.exp(-np.dot(X_test,mm)*kappa))
P = np.reshape(p_test, X1.shape)
CS = plt.contour(X1, X2, P, [0.1,0.175,0.25,0.375, 0.5,0.625,0.75,0.825,0.9])
plt.clabel(CS)
plt.xlabel('x_1')
plt.ylabel('x_2')
plt.title('Contours of p(y=1|x,D)')
plt.show()
```

```python
torch.nn.Softplus()
```

    Softplus(beta=1, threshold=20)

## MAP Baseline

We run the point estimate approximation (Maximum a posteriori) to double check what the learned weights look like.  We get the  exact same training accuracy as with the controlled model and similarly large weights for the non bias weights.

```python
X.shape
```

    (100, 3)

```python
Θ_map = torch.zeros((1, dim), requires_grad=True, device=device)
optimizer_map = torch.optim.Adam([Θ_map], lr=0.05)
#     optimizer = torch.optim.LBFGS(gpr.parameters(), lr=0.01)

losses_map = []
num_steps = 1000
for i in tqdm(range(num_steps)):
    optimizer_map.zero_grad()

    if isinstance(optimizer_map, torch.optim.LBFGS):
        def closure_map():
            loss_map = log_likelihood_vmap()
            optimizer_map.zero_grad()
            loss_map.backward()
            return loss

        optimizer_map.step(closure_map)
        losses_map.append(closure_map().item())
    else:
        loss_map = -(log_likelihood_vmap(Θ_map, X_train, y_train) + gaussian_prior(Θ_map))
        optimizer_map.zero_grad()
        loss_map.backward()
        print(loss_map.item())
        optimizer_map.step()
        losses_map.append(loss_map.item())

Θ_map
pred_map = torch.sigmoid(X_train.mm(Θ_map.T)).mean(axis=1)
((pred_map < 0.5).float() == y_train).float().mean(), Θ_map
```

      0%|          | 0/1000 [00:00<?, ?it/s]

    C:\Users\vargf\AppData\Local\Temp/ipykernel_2136/381526447.py:36: UserWarning: torch.vmap is an experimental prototype that is subject to change and/or deletion. Please use at your own risk. There may be unexpected performance cliffs due to certain operators not being implemented. To see detailed performance warnings please use `torch._C._debug_only_display_vmap_fallback_warnings(True) before the call to `vmap`.
      batched_loss =  torch._vmap_internals.vmap(loss_)

    55.45177459716797
    52.67886734008789
    50.03776931762695
    47.52959442138672
    45.154483795166016
    42.91158676147461
    40.79913330078125
    38.81449508666992
    36.95428466796875
    35.21445083618164
    33.590423583984375
    32.07719421386719
    30.66946029663086
    29.36173439025879
    28.148420333862305
    27.023902893066406
    25.982643127441406
    25.019207000732422
    24.128328323364258
    23.30495262145996
    22.54424476623535
    21.841611862182617
    21.192726135253906
    20.593515396118164
    20.040164947509766
    19.529117584228516
    19.057064056396484
    18.620927810668945
    18.21786880493164
    17.845247268676758
    17.500642776489258
    17.181812286376953
    16.886699676513672
    16.613401412963867
    16.360179901123047
    16.12543487548828
    15.907696723937988
    15.705615997314453
    15.517951011657715
    15.343571662902832
    15.181438446044922
    15.030595779418945
    14.890161514282227
    14.759340286254883
    14.637393951416016
    14.523643493652344
    14.417470932006836
    14.318305969238281
    14.225625991821289
    14.13895034790039
    14.057838439941406
    13.981886863708496
    13.910720825195312
    13.843996047973633
    13.781397819519043
    13.722637176513672
    13.66744327545166
    13.61557388305664
    13.566798210144043
    13.520904541015625
    13.477703094482422
    13.437012672424316
    13.398666381835938
    13.362512588500977
    13.328409194946289
    13.296223640441895
    13.265838623046875
    13.23713493347168
    13.210012435913086
    13.184370040893555
    13.160122871398926
    13.13718318939209
    13.115469932556152
    13.094917297363281
    13.07545280456543
    13.057011604309082
    13.039539337158203
    13.022977828979492
    13.00727653503418
    12.99238395690918
    12.978260040283203
    12.964859008789062
    12.952140808105469
    12.940072059631348
    12.92861557006836
    12.91773796081543
    12.907407760620117
    12.897597312927246
    12.88828182220459
    12.879429817199707
    12.87102222442627
    12.863033294677734
    12.855443000793457
    12.848231315612793
    12.841375350952148
    12.834863662719727
    12.8286714553833
    12.822789192199707
    12.81719970703125
    12.81188678741455
    12.80683708190918
    12.802037239074707
    12.797475814819336
    12.793142318725586
    12.78902530670166
    12.785109519958496
    12.781390190124512
    12.777854919433594
    12.774497032165527
    12.771309852600098
    12.768280029296875
    12.765399932861328
    12.762666702270508
    12.76007080078125
    12.757606506347656
    12.755267143249512
    12.753044128417969
    12.750932693481445
    12.748931884765625
    12.747032165527344
    12.745229721069336
    12.743517875671387
    12.741894721984863
    12.740355491638184
    12.738895416259766
    12.737510681152344
    12.736200332641602
    12.734954833984375
    12.73377513885498
    12.732657432556152
    12.731599807739258
    12.730597496032715
    12.729646682739258
    12.72874641418457
    12.727895736694336
    12.727088928222656
    12.726325988769531
    12.725605964660645
    12.724923133850098
    12.72427749633789
    12.723665237426758
    12.723087310791016
    12.722541809082031
    12.722025871276855
    12.721537590026855
    12.721076965332031
    12.720643997192383
    12.720233917236328
    12.719844818115234
    12.719480514526367
    12.719133377075195
    12.71881103515625
    12.718502044677734
    12.718212127685547
    12.717939376831055
    12.717683792114258
    12.717440605163574
    12.717211723327637
    12.716997146606445
    12.71679401397705
    12.716605186462402
    12.716425895690918
    12.716257095336914
    12.71609878540039
    12.715950012207031
    12.715808868408203
    12.715678215026855
    12.715555191040039
    12.715438842773438
    12.715330123901367
    12.715227127075195
    12.715131759643555
    12.715041160583496
    12.714957237243652
    12.71487808227539
    12.714803695678711
    12.714734077453613
    12.714668273925781
    12.714609146118164
    12.71455192565918
    12.714496612548828
    12.714447021484375
    12.714401245117188
    12.714357376098633
    12.714315414428711
    12.714277267456055
    12.714242935180664
    12.714208602905273
    12.714178085327148
    12.714149475097656
    12.714122772216797
    12.714097023010254
    12.714075088500977
    12.7140531539917
    12.714031219482422
    12.71401309967041
    12.713996887207031
    12.713979721069336
    12.713964462280273
    12.713949203491211
    12.713935852050781
    12.713923454284668
    12.713912963867188
    12.71390151977539
    12.713892936706543
    12.713883399963379
    12.713874816894531
    12.7138671875
    12.713861465454102
    12.71385383605957
    12.713848114013672
    12.71384048461914
    12.713836669921875
    12.713830947875977
    12.713824272155762
    12.713821411132812
    12.713817596435547
    12.713814735412598
    12.713810920715332
    12.713809967041016
    12.713804244995117
    12.7138032913208
    12.713800430297852
    12.713799476623535
    12.713797569274902
    12.713794708251953
    12.71379280090332
    12.713790893554688
    12.713790893554688
    12.713788986206055
    12.713787078857422
    12.713786125183105
    12.713785171508789
    12.713784217834473
    12.713783264160156
    12.713784217834473
    12.713781356811523
    12.713783264160156
    12.713780403137207
    12.71377944946289
    12.713780403137207
    12.713780403137207
    12.71377944946289
    12.713780403137207
    12.713778495788574
    12.713778495788574
    12.713778495788574
    12.713777542114258
    12.713777542114258
    12.713775634765625
    12.713775634765625
    12.713777542114258
    12.713776588439941
    12.713777542114258
    12.713776588439941
    12.713777542114258
    12.713777542114258
    12.713774681091309
    12.713775634765625
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713774681091309
    12.713775634765625
    12.713775634765625
    12.713773727416992
    12.713774681091309
    12.713775634765625
    12.713775634765625
    12.713776588439941
    12.713773727416992
    12.713776588439941
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713774681091309
    12.713775634765625
    12.713776588439941
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713776588439941
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713776588439941
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713774681091309
    12.713774681091309
    12.713774681091309
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713773727416992
    12.713773727416992
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713773727416992
    12.713774681091309
    12.713773727416992
    12.713773727416992
    12.713774681091309
    12.713773727416992
    12.713775634765625
    12.713775634765625
    12.713774681091309
    12.713774681091309
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713774681091309
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713774681091309
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713776588439941
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625
    12.713775634765625

    (tensor(0., device='cuda:0'),
     tensor([[ 0.0000,  2.4463, -2.1438]], device='cuda:0', requires_grad=True))

## Pyro Bayesian Logistic Regeression

As a baseline we run pyro with SVI on this same example atm we get very different results which is a bit worrying.

#### EDIT:

I think results are actually in agreement what was happening is the priors in the pyro code where much more confident thus the smaller weights. I just tried changing the prior of this pyro model and the posterior is not changing much, this is a bit fishy maybe worth investigating.

I think it might just be the case that this model is not very good.

```python
import numpy as np
import scipy.special as ssp
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.distributions.constraints as constraints

from torch.utils.data import DataLoader
from torch.utils.data.sampler import SubsetRandomSampler

import pyro
import pyro.distributions as dist

from pyro.infer import SVI, Trace_ELBO
from pyro.optim import Adam, SGD

pyro.enable_validation(True)
torch.set_default_dtype(torch.double) # this was necessary on the CPU
```

```python

# these were adapted from the Pyro VAE tutorial

def train(svi, train_loader, n_train):
    # initialize loss accumulator
    epoch_loss = 0.
    # do a training epoch over each mini-batch x returned
    # by the data loader
    for _, xs in enumerate(train_loader):
        # do ELBO gradient and accumulate loss
        epoch_loss += svi.step(*xs)

    # return epoch loss
    total_epoch_loss_train = epoch_loss / n_train
    return total_epoch_loss_train

def evaluate(svi, test_loader, n_test):
    # initialize loss accumulator
    test_loss = 0.
    # compute the loss over the entire test set
    for _, xs in enumerate(test_loader):
        # compute ELBO estimate and accumulate loss
        test_loss += svi.evaluate_loss(*xs)

    total_epoch_loss_test = test_loss / n_test
    return total_epoch_loss_test

def plot_llk(train_elbo, test_elbo, test_int):
    plt.figure(figsize=(8, 6))

    x = np.arange(len(train_elbo))

    plt.plot(x, train_elbo, marker='o', label='Train ELBO')
    plt.plot(x[::test_int], test_elbo, marker='o', label='Test ELBO')
    plt.xlabel('Training Epoch')
    plt.legend()
    plt.show()
```

```python
class LogRegressionModel(nn.Module):
    def __init__(self, p):
        super(LogRegressionModel, self).__init__()

        self.p = p

        # hyperparameters for normal priors
        self.alpha_h_loc = torch.zeros(1, p)
        self.alpha_h_scale = 0.0001 * torch.ones(1, p)
        self.beta_h_loc = torch.zeros(1)
        self.beta_h_scale =  0.0001 * torch.ones(1)

        # initial values of variational parameters
        self.alpha_0 = np.zeros((1, p))
        self.alpha_0_scale = np.ones((1, p))
        self.beta_0 = np.zeros((1,))
        self.beta_0_scale = np.ones((1,))

    def model(self, x, y):
        # sample from prior
        a = pyro.sample(
            "weight", dist.Normal(self.alpha_h_loc, self.alpha_h_scale, validate_args=True).independent(1)
        ).float()
        b = pyro.sample(
            "bias", dist.Normal(self.beta_h_loc, self.beta_h_scale, validate_args=True).independent(1)
        ).float()

        with pyro.iarange("data", x.size(0)):
#             import pdb; pdb.set_trace()
            model_logits = (torch.matmul(x, a.permute(1, 0)) + b).squeeze()

            pyro.sample(
                "obs",
                dist.Bernoulli(logits=model_logits, validate_args=True),
                obs=y.squeeze()
            )

    def guide(self, x, y):
        # register variational parameters with pyro
        alpha_loc = pyro.param("alpha_loc", torch.tensor(self.alpha_0))
        alpha_scale = pyro.param("alpha_scale", torch.tensor(self.alpha_0_scale),
                                 constraint=constraints.positive)
        beta_loc = pyro.param("beta_loc", torch.tensor(self.beta_0))
        beta_scale = pyro.param("beta_scale", torch.tensor(self.beta_0_scale),
                                constraint=constraints.positive)

        pyro.sample(
            "weight", dist.Normal(alpha_loc, alpha_scale, validate_args=True).independent(1)
        )
        pyro.sample(
            "bias", dist.Normal(beta_loc, beta_scale, validate_args=True).independent(1)
        )

    def guide_(self, x, y, N=100):
        # register variational parameters with pyro
        alpha_loc = pyro.param("alpha_loc", torch.tensor(self.alpha_0))
        alpha_scale = pyro.param("alpha_scale", torch.tensor(self.alpha_0_scale),
                                 constraint=constraints.positive)
        beta_loc = pyro.param("beta_loc", torch.tensor(self.beta_0))
        beta_scale = pyro.param("beta_scale", torch.tensor(self.beta_0_scale),
                                constraint=constraints.positive)
        w = []
        b = []
        for _ in range(N):
            w.append(torch.tensor(pyro.sample(
                "weight", dist.Normal(alpha_loc, alpha_scale, validate_args=True).independent(1)
            )))
            b.append(torch.tensor(pyro.sample(
                "bias", dist.Normal(beta_loc, beta_scale, validate_args=True).independent(1)
            )))
#         import pdb;pdb.set_trace()
        return torch.vstack(w), torch.vstack(b)
```

```python

pyro.clear_param_store()

optim = Adam({'lr': 0.01})

num_epochs = 1000
batch_size = 50

N = X.shape[0]
p = 2

n_train = X_train.shape[0]
example_indices = np.random.permutation(N)
n_test = N - n_train
test_iter = 50
```

```python
lr_model = LogRegressionModel(p=p)

svi = SVI(
    lr_model.model, lr_model.guide, optim,
    loss=Trace_ELBO()
)

lr_dataset = torch.utils.data.TensorDataset(torch.tensor(X[:,1:]).float(), torch.tensor(y.reshape(-1,1)).float())

data_loader_train = DataLoader(
    dataset=lr_dataset, batch_size=batch_size, pin_memory=False,
    sampler=SubsetRandomSampler(example_indices[:n_train]),
)

data_loader_test = DataLoader(
    dataset=lr_dataset, batch_size=batch_size, pin_memory=False,
    sampler=SubsetRandomSampler(example_indices[n_train:]),
)

train_elbo = []
test_elbo = []
for epoch in range(num_epochs):
    total_epoch_loss_train = train(svi, data_loader_train, n_train)
    train_elbo.append(-total_epoch_loss_train)

    if epoch % test_iter == 0:
        print("[epoch %03d]  average training loss: %.4f" % (epoch, total_epoch_loss_train))
        # report test diagnostics
        total_epoch_loss_test = evaluate(svi, data_loader_test, n_test)
        test_elbo.append(-total_epoch_loss_test)
```

    [epoch 000]  average training loss: 1220005.2042
    [epoch 050]  average training loss: 316690.8360
    [epoch 100]  average training loss: 578170.2876
    [epoch 150]  average training loss: 568008.1648
    [epoch 200]  average training loss: 307361.8592
    [epoch 250]  average training loss: 131007.7829
    [epoch 300]  average training loss: 73617.5958
    [epoch 350]  average training loss: 191313.3133
    [epoch 400]  average training loss: 111897.0367
    [epoch 450]  average training loss: 300512.3689
    [epoch 500]  average training loss: 19642.3765
    [epoch 550]  average training loss: 59578.3530
    [epoch 600]  average training loss: 97903.6400
    [epoch 650]  average training loss: 55632.0432
    [epoch 700]  average training loss: 116581.4045
    [epoch 750]  average training loss: 31808.7584
    [epoch 800]  average training loss: 91703.4840
    [epoch 850]  average training loss: 20083.9464
    [epoch 900]  average training loss: 16839.9660
    [epoch 950]  average training loss: 48150.2470

```python
plot_llk(train_elbo, test_elbo, test_iter)
```

```python
W, b = lr_model.guide_(X_train.float(), y_train.float())
```

    C:\Users\vargf\AppData\Local\Temp/ipykernel_17988/447687565.py:66: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
      w.append(torch.tensor(pyro.sample(
    C:\Users\vargf\AppData\Local\Temp/ipykernel_17988/447687565.py:69: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
      b.append(torch.tensor(pyro.sample(

```python
fig, (ax1,ax2,ax3) = plt.subplots(1,3)

# ax1.hist(Θ_1[:,0].detach().numpy())
# ax2.hist(Θ_1[:,1].detach().numpy())
# ax3.hist(Θ_1[:,2].detach().numpy())

ax1.hist(b.detach().numpy(), color="red")
ax2.hist(W[:,0].detach().numpy(), color="red")
ax3.hist(W[:,1].detach().numpy(), color="red")

```

```python
b.mean(), W.mean(axis=0)
```

    (tensor(-0.0138), tensor([0.0121, 0.0196]))

```python

```