debug.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import hydra\n",
    "from omegaconf import DictConfig, OmegaConf\n",
    "import sys\n",
    "import torch\n",
    "import numpy as np\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "torch.backends.cudnn.benchmark = True\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bayesian_dpddm import ConvModel, DPDDMBayesianMonitor\n",
    "from experiments.utils import get_datasets, get_configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x77f89c3a5e10>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "RANDOM_SEED = 9927\n",
    "np.random.seed(RANDOM_SEED)\n",
    "torch.random.manual_seed(RANDOM_SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get Hydra configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_dir = 'experiments/configs/'\n",
    "config_name = 'cifar10'\n",
    "hydra.initialize(config_path=config_dir, version_base='1.2')\n",
    "cfg = hydra.compose(config_name=config_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg.from_pretrained = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset =  get_datasets(cfg)\n",
    "model_config, train_config = get_configs(cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build model and monitor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_model = ConvModel(model_config, train_size=len(dataset['train']))\n",
    "monitor = DPDDMBayesianMonitor(\n",
    "    model=base_model,\n",
    "    trainset=dataset['train'],\n",
    "    valset=dataset['valid'],\n",
    "    train_cfg=train_config,\n",
    "    device=device,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "base_model.load_state_dict(torch.load(os.path.join('saved_weights', f'{cfg.dataset.name}.pth')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sanity check trained performance\n",
    "batch_size = 128\n",
    "train_loader = DataLoader(dataset['train'], batch_size=batch_size)\n",
    "val_loader = DataLoader(dataset['valid'], batch_size=batch_size)\n",
    "test_loader = DataLoader(dataset['dpddm_id'], batch_size=batch_size)\n",
    "ood_loader = DataLoader(dataset['dpddm_ood'], batch_size=batch_size)\n",
    "loaders = [train_loader, val_loader, test_loader, ood_loader]\n",
    "loaders_id = ['train', 'valid', 'test', 'ood']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loader: train, Accuracy: 0.865390375399361\n",
      "Loader: valid, Accuracy: 0.8375197784810127\n",
      "Loader: test, Accuracy: 0.8306962025316456\n",
      "Loader: ood, Accuracy: 0.7216796875\n"
     ]
    }
   ],
   "source": [
    "for loader, loader_id in zip(loaders, loaders_id):\n",
    "    accs = [] \n",
    "    monitor.model.eval()\n",
    "    with torch.no_grad():\n",
    "        for images, labels in loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            output = monitor.model(images)\n",
    "            probs = output.predictive.probs\n",
    "            acc = monitor.eval_acc(probs, labels).item()\n",
    "            accs.append(acc)\n",
    "    print(f\"Loader: {loader_id}, Accuracy: {np.mean(accs)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 23%|██▎       | 114/500 [00:16<00:54,  7.11it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[11], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m args \u001b[38;5;241m=\u001b[39m cfg\n\u001b[0;32m----> 2\u001b[0m monitor\u001b[38;5;241m.\u001b[39mpretrain_disagreement_distribution(dataset\u001b[38;5;241m=\u001b[39mdataset[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdpddm_train\u001b[39m\u001b[38;5;124m'\u001b[39m],\n\u001b[1;32m      3\u001b[0m                                                n_post_samples\u001b[38;5;241m=\u001b[39margs\u001b[38;5;241m.\u001b[39mdpddm\u001b[38;5;241m.\u001b[39mn_post_samples,\n\u001b[1;32m      4\u001b[0m                                                data_sample_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m500\u001b[39m,\n\u001b[1;32m      5\u001b[0m                                                Phi_size\u001b[38;5;241m=\u001b[39margs\u001b[38;5;241m.\u001b[39mdpddm\u001b[38;5;241m.\u001b[39mPhi_size, \n\u001b[1;32m      6\u001b[0m                                                temperature\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m,\n\u001b[1;32m      7\u001b[0m                                                )\n",
      "File \u001b[0;32m~/bayesian_dpddm/bayesian_dpddm/monitors/monitor.py:127\u001b[0m, in \u001b[0;36mDPDDMMonitor.pretrain_disagreement_distribution\u001b[0;34m(self, dataset, Phi_size, tqdm_enabled, *args, **kwargs)\u001b[0m\n\u001b[1;32m    125\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[1;32m    126\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m f(\u001b[38;5;28mrange\u001b[39m(Phi_size)):\n\u001b[0;32m--> 127\u001b[0m         max_dis_rate, _ \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdpddm_test(dataset, replace\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreplace, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    128\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mPhi\u001b[38;5;241m.\u001b[39mappend(max_dis_rate)\n",
      "File \u001b[0;32m~/bayesian_dpddm/bayesian_dpddm/monitors/monitor.py:109\u001b[0m, in \u001b[0;36mDPDDMMonitor.dpddm_test\u001b[0;34m(self, dataset, data_sample_size, alpha, replace, *args, **kwargs)\u001b[0m\n\u001b[1;32m     93\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Given a dataset, computes the maximum disagreement rate as well as the OOD verdict.\u001b[39;00m\n\u001b[1;32m     94\u001b[0m \u001b[38;5;124;03mUsed to both generate Phi and Algorithms 2 and 4.\u001b[39;00m\n\u001b[1;32m     95\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    106\u001b[0m \u001b[38;5;124;03m    - OOD verdict w.r.t. self.Phi\u001b[39;00m\n\u001b[1;32m    107\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    108\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[0;32m--> 109\u001b[0m     X, _ \u001b[38;5;241m=\u001b[39m sample_from_dataset(n_samples\u001b[38;5;241m=\u001b[39mdata_sample_size, dataset\u001b[38;5;241m=\u001b[39mdataset, replace\u001b[38;5;241m=\u001b[39mreplace)\n\u001b[1;32m    110\u001b[0m     y_pseudo \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_pseudolabels(X)\n\u001b[1;32m    111\u001b[0m     max_dis_rate \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcompute_max_dis_rate(X, y_pseudo, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[0;32m~/bayesian_dpddm/bayesian_dpddm/monitors/utils.py:27\u001b[0m, in \u001b[0;36msample_from_dataset\u001b[0;34m(n_samples, dataset, replace)\u001b[0m\n\u001b[1;32m     25\u001b[0m true_labels \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mzeros(size\u001b[38;5;241m=\u001b[39m(n_samples,))\n\u001b[1;32m     26\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mlen\u001b[39m(indices)):\n\u001b[0;32m---> 27\u001b[0m     tmp[i] \u001b[38;5;241m=\u001b[39m dataset[indices[i]][\u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m     28\u001b[0m     true_labels[i] \u001b[38;5;241m=\u001b[39m dataset[indices[i]][\u001b[38;5;241m1\u001b[39m]\n\u001b[1;32m     29\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m tmp, true_labels\n",
      "File \u001b[0;32m~/anaconda3/envs/dpddm/lib/python3.11/site-packages/torch/utils/data/dataset.py:412\u001b[0m, in \u001b[0;36mSubset.__getitem__\u001b[0;34m(self, idx)\u001b[0m\n\u001b[1;32m    410\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(idx, \u001b[38;5;28mlist\u001b[39m):\n\u001b[1;32m    411\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mindices[i] \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m idx]]\n\u001b[0;32m--> 412\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mindices[idx]]\n",
      "File \u001b[0;32m~/anaconda3/envs/dpddm/lib/python3.11/site-packages/torchvision/datasets/cifar.py:116\u001b[0m, in \u001b[0;36mCIFAR10.__getitem__\u001b[0;34m(self, index)\u001b[0m\n\u001b[1;32m    112\u001b[0m img, target \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdata[index], \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtargets[index]\n\u001b[1;32m    114\u001b[0m \u001b[38;5;66;03m# doing this so that it is consistent with all other datasets\u001b[39;00m\n\u001b[1;32m    115\u001b[0m \u001b[38;5;66;03m# to return a PIL Image\u001b[39;00m\n\u001b[0;32m--> 116\u001b[0m img \u001b[38;5;241m=\u001b[39m Image\u001b[38;5;241m.\u001b[39mfromarray(img)\n\u001b[1;32m    118\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtransform \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    119\u001b[0m     img \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtransform(img)\n",
      "File \u001b[0;32m~/anaconda3/envs/dpddm/lib/python3.11/site-packages/PIL/Image.py:3331\u001b[0m, in \u001b[0;36mfromarray\u001b[0;34m(obj, mode)\u001b[0m\n\u001b[1;32m   3329\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m strides \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m   3330\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(obj, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtobytes\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m-> 3331\u001b[0m         obj \u001b[38;5;241m=\u001b[39m obj\u001b[38;5;241m.\u001b[39mtobytes()\n\u001b[1;32m   3332\u001b[0m     \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(obj, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtostring\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[1;32m   3333\u001b[0m         obj \u001b[38;5;241m=\u001b[39m obj\u001b[38;5;241m.\u001b[39mtostring()\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "args = cfg\n",
    "monitor.pretrain_disagreement_distribution(dataset=dataset['dpddm_train'],\n",
    "                                               n_post_samples=args.dpddm.n_post_samples,\n",
    "                                               data_sample_size=500,\n",
    "                                               Phi_size=args.dpddm.Phi_size, \n",
    "                                               temperature=10,\n",
    "                                               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:01<00:00, 81.47it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "np.float64(0.98)"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rate, max_dis_rates = monitor.repeat_tests(n_repeats=args.dpddm.n_repeats,\n",
    "                                      dataset=dataset['dpddm_train'], \n",
    "                                      n_post_samples=args.dpddm.n_post_samples,\n",
    "                                      data_sample_size=50,\n",
    "                                      temperature=10\n",
    "                                      )\n",
    "rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 2.,  7., 26., 26.,  0., 17., 14.,  7.,  0.,  1.]),\n",
       " array([0.08      , 0.096     , 0.112     , 0.128     , 0.144     ,\n",
       "        0.16      , 0.176     , 0.192     , 0.208     , 0.22399999,\n",
       "        0.23999999]),\n",
       " <BarContainer object of 10 artists>)"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAigAAAGdCAYAAAA44ojeAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQAAIKJJREFUeJzt3X1s1eX9//HXsXdA055RCj3tqKVbUKLt3CwK1BtQoNiARDEDRmbAoMEIjK4QVjQLdVtaZBuYDWUZIdw5bpyKmkCEGqXKOiZ2GOUmDgfMEnvSwfC0RXYKeH3/2I/z89CinHJOz7vt85GchPM5V8/nus7pxz799JxTj3POCQAAwJDr4j0BAACAyxEoAADAHAIFAACYQ6AAAABzCBQAAGAOgQIAAMwhUAAAgDkECgAAMCcx3hPojC+//FKfffaZ0tLS5PF44j0dAABwFZxzamlpUU5Ojq677uvPkXTLQPnss8+Um5sb72kAAIBOaGho0ODBg792TLcMlLS0NEn/W2B6enqcZwMAAK5Gc3OzcnNzQz/Hv063DJRLv9ZJT08nUAAA6Gau5uUZvEgWAACYQ6AAAABzCBQAAGAOgQIAAMwhUAAAgDkECgAAMIdAAQAA5hAoAADAHAIFAACYQ6AAAABzCBQAAGAOgQIAAMwhUAAAgDkECgAAMCcx3hNAx4ZU7LjqsSeWTYzhTAAA6HqcQQEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwJyIAqW6ulq33Xab0tLSNGjQID3wwAP6+OOPw8bMmjVLHo8n7DJy5MiwMcFgUPPnz1dmZqZSU1M1efJknTx58tpXAwAAeoSIAqW2tlZz587Vvn37VFNTowsXLqikpERnz54NG3ffffepsbExdNm5c2fY7WVlZdq+fbu2bt2qvXv3qrW1VZMmTdLFixevfUUAAKDbS4xk8BtvvBF2fd26dRo0aJDq6+t19913h7anpKTI5/N1eB+BQEBr167Vpk2bNG7cOEnSCy+8oNzcXL355puaMGFCpGsAAAA9zDW9BiUQCEiSMjIywrbv2bNHgwYN0g033KDHHntMTU1Nodvq6+t1/vx5lZSUhLbl5OSooKBAdXV1He4nGAyqubk57AIAAHquTgeKc07l5eW68847VVBQENpeWlqqP/3pT3rrrbf029/+Vvv379e9996rYDAoSfL7/UpOTlb//v3D7i8rK0t+v7/DfVVXV8vr9YYuubm5nZ02AADoBiL6Fc9XzZs3Tx9++KH27t0btn3atGmhfxcUFGj48OHKy8vTjh07NGXKlCven3NOHo+nw9uWLFmi8vLy0PXm5mYiBQCAHqxTZ1Dmz5+v119/XW+//bYGDx78tWOzs7OVl5eno0ePSpJ8Pp/a2tp05syZsHFNTU3Kysrq8D5SUlKUnp4edgEAAD1XRIHinNO8efP0yiuv6K233lJ+fv43fs3p06fV0NCg7OxsSVJRUZGSkpJUU1MTGtPY2KiDBw+quLg4wukDAICeKKJf8cydO1ebN2/Wa6+9prS0tNBrRrxer/r27avW1lZVVlbqoYceUnZ2tk6cOKEnn3xSmZmZevDBB0NjZ8+erYULF2rAgAHKyMjQokWLVFhYGHpXDwAA6N0iCpTVq1dLksaMGRO2fd26dZo1a5YSEhL00UcfaePGjfr888+VnZ2te+65R9u2bVNaWlpo/MqVK5WYmKipU6fq3LlzGjt2rNavX6+EhIRrXxEAAOj2PM45F+9JRKq5uVler1eBQKDHvh5lSMWOqx57YtnEGM4EAIDoiOTnN3+LBwAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5EQVKdXW1brvtNqWlpWnQoEF64IEH9PHHH4eNcc6psrJSOTk56tu3r8aMGaNDhw6FjQkGg5o/f74yMzOVmpqqyZMn6+TJk9e+GgAA0CNEFCi1tbWaO3eu9u3bp5qaGl24cEElJSU6e/ZsaMzy5cu1YsUKrVq1Svv375fP59P48ePV0tISGlNWVqbt27dr69at2rt3r1pbWzVp0iRdvHgxeisDAADdlsc55zr7xf/+9781aNAg1dbW6u6775ZzTjk5OSorK9PPfvYzSf87W5KVlaVnnnlGc+bMUSAQ0MCBA7Vp0yZNmzZNkvTZZ58pNzdXO3fu1IQJE75xv83NzfJ6vQoEAkpPT+/s9E0bUrHjqseeWDYxhjMBACA6Ivn5fU2vQQkEApKkjIwMSdLx48fl9/tVUlISGpOSkqLRo0errq5OklRfX6/z58+HjcnJyVFBQUFozOWCwaCam5vDLgAAoOfqdKA451ReXq4777xTBQUFkiS/3y9JysrKChublZUVus3v9ys5OVn9+/e/4pjLVVdXy+v1hi65ubmdnTYAAOgGOh0o8+bN04cffqgtW7a0u83j8YRdd86123a5rxuzZMkSBQKB0KWhoaGz0wYAAN1ApwJl/vz5ev311/X2229r8ODBoe0+n0+S2p0JaWpqCp1V8fl8amtr05kzZ6445nIpKSlKT08PuwAAgJ4rokBxzmnevHl65ZVX9NZbbyk/Pz/s9vz8fPl8PtXU1IS2tbW1qba2VsXFxZKkoqIiJSUlhY1pbGzUwYMHQ2MAAEDvlhjJ4Llz52rz5s167bXXlJaWFjpT4vV61bdvX3k8HpWVlamqqkpDhw7V0KFDVVVVpX79+mnGjBmhsbNnz9bChQs1YMAAZWRkaNGiRSosLNS4ceOiv0IAANDtRBQoq1evliSNGTMmbPu6des0a9YsSdLixYt17tw5PfHEEzpz5oxGjBih3bt3Ky0tLTR+5cqVSkxM1NSpU3Xu3DmNHTtW69evV0JCwrWtBgAA9AjX9Dko8cLnoITjc1AAAN1Bl30OCgAAQCwQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADAnMd4T6E2GVOyI9xQAAOgWOIMCAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcxLjPQFcuyEVOyIaf2LZxBjNBACA6OAMCgAAMIdAAQAA5hAoAADAHAIFAACYQ6AAAABzCBQAAGAOgQIAAMwhUAAAgDkECgAAMIdAAQAA5hAoAADAHAIFAACYE3GgvPPOO7r//vuVk5Mjj8ejV199Nez2WbNmyePxhF1GjhwZNiYYDGr+/PnKzMxUamqqJk+erJMnT17TQgAAQM8RcaCcPXtWt9xyi1atWnXFMffdd58aGxtDl507d4bdXlZWpu3bt2vr1q3au3evWltbNWnSJF28eDHyFQAAgB4nMdIvKC0tVWlp6deOSUlJkc/n6/C2QCCgtWvXatOmTRo3bpwk6YUXXlBubq7efPNNTZgwIdIpAQCAHiYmr0HZs2ePBg0apBtuuEGPPfaYmpqaQrfV19fr/PnzKikpCW3LyclRQUGB6urqYjEdAADQzUR8BuWblJaW6oc//KHy8vJ0/Phx/fznP9e9996r+vp6paSkyO/3Kzk5Wf379w/7uqysLPn9/g7vMxgMKhgMhq43NzdHe9oAAMCQqAfKtGnTQv8uKCjQ8OHDlZeXpx07dmjKlClX/DrnnDweT4e3VVdX6+mnn472VAEAgFExf5txdna28vLydPToUUmSz+dTW1ubzpw5EzauqalJWVlZHd7HkiVLFAgEQpeGhoZYTxsAAMRRzAPl9OnTamhoUHZ2tiSpqKhISUlJqqmpCY1pbGzUwYMHVVxc3OF9pKSkKD09PewCAAB6roh/xdPa2qpPPvkkdP348eP64IMPlJGRoYyMDFVWVuqhhx5Sdna2Tpw4oSeffFKZmZl68MEHJUler1ezZ8/WwoULNWDAAGVkZGjRokUqLCwMvasHAAD0bhEHyvvvv6977rkndL28vFySNHPmTK1evVofffSRNm7cqM8//1zZ2dm65557tG3bNqWlpYW+ZuXKlUpMTNTUqVN17tw5jR07VuvXr1dCQkIUlgQAALo7j3POxXsSkWpubpbX61UgEOhWv+4ZUrEj3lOQJJ1YNjHeUwAA9EKR/Pzmb/EAAABzCBQAAGAOgQIAAMwhUAAAgDkECgAAMIdAAQAA5hAoAADAHAIFAACYQ6AAAABzCBQAAGAOgQIAAMwhUAAAgDkECgAAMIdAAQAA5hAoAADAHAIFAACYQ6AAAABzCBQAAGAOgQIAAMwhUAAAgDkECgAAMCcx3hNA1xtSseOqx55YNjGGMwEAoGOcQQEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADAnMd4TgG1DKnZc9dgTyybGcCYAgN6EMygAAMAcAgUAAJhDoAAAAHMIFAAAYA6BAgAAzCFQAACAOQQKAAAwJ+JAeeedd3T//fcrJydHHo9Hr776atjtzjlVVlYqJydHffv21ZgxY3To0KGwMcFgUPPnz1dmZqZSU1M1efJknTx58poWAgAAeo6IA+Xs2bO65ZZbtGrVqg5vX758uVasWKFVq1Zp//798vl8Gj9+vFpaWkJjysrKtH37dm3dulV79+5Va2urJk2apIsXL3Z+JQAAoMeI+JNkS0tLVVpa2uFtzjk9++yzeuqppzRlyhRJ0oYNG5SVlaXNmzdrzpw5CgQCWrt2rTZt2qRx48ZJkl544QXl5ubqzTff1IQJE65hOQAAoCeI6kfdHz9+XH6/XyUlJaFtKSkpGj16tOrq6jRnzhzV19fr/PnzYWNycnJUUFCgurq6DgMlGAwqGAyGrjc3N0dz2tckko+CBwAAVyeqL5L1+/2SpKysrLDtWVlZodv8fr+Sk5PVv3//K465XHV1tbxeb+iSm5sbzWkDAABjYvIuHo/HE3bdOddu2+W+bsySJUsUCARCl4aGhqjNFQAA2BPVQPH5fJLU7kxIU1NT6KyKz+dTW1ubzpw5c8Uxl0tJSVF6enrYBQAA9FxRDZT8/Hz5fD7V1NSEtrW1tam2tlbFxcWSpKKiIiUlJYWNaWxs1MGDB0NjAABA7xbxi2RbW1v1ySefhK4fP35cH3zwgTIyMnT99derrKxMVVVVGjp0qIYOHaqqqir169dPM2bMkCR5vV7Nnj1bCxcu1IABA5SRkaFFixapsLAw9K4eAADQu0UcKO+//77uueee0PXy8nJJ0syZM7V+/XotXrxY586d0xNPPKEzZ85oxIgR2r17t9LS0kJfs3LlSiUmJmrq1Kk6d+6cxo4dq/Xr1yshISEKSwIAAN2dxznn4j2JSDU3N8vr9SoQCMT99Si8zfj/O7FsYrynAAAwLJKf3/wtHgAAYA6BAgAAzCFQAACAOQQKAAAwh0ABAADmECgAAMAcAgUAAJhDoAAAAHMIFAAAYA6BAgAAzCFQAACAOQQKAAAwh0ABAADmECgAAMAcAgUAAJhDoAAAAHMIFAAAYA6BAgAAzCFQAACAOQQKAAAwh0ABAADmECgAAMAcAgUAAJhDoAAAAHMIFAAAYA6BAgAAzCFQAACAOQQKAAAwh0ABAADmECgAAMAcAgUAAJhDoAAAAHMIFAAAYA6BAgAAzCFQAACAOQQKAAAwh0ABAADmECgAAMAcAgUAAJhDoAAAAHMIFAAAYA6BAgAAzCFQAACAOQQKAAAwh0ABAADmECgAAMAcAgUAAJhDoAAAAHMIFAAAYA6BAgAAzCFQAACAOQQKAAAwh0ABAADmRD1QKisr5fF4wi4+ny90u3NOlZWVysnJUd++fTVmzBgdOnQo2tMAAADdWEzOoNx8881qbGwMXT766KPQbcuXL9eKFSu0atUq7d+/Xz6fT+PHj1dLS0sspgIAALqhmARKYmKifD5f6DJw4EBJ/zt78uyzz+qpp57SlClTVFBQoA0bNuiLL77Q5s2bYzEVAADQDcUkUI4ePaqcnBzl5+dr+vTpOnbsmCTp+PHj8vv9KikpCY1NSUnR6NGjVVdXd8X7CwaDam5uDrsAAICeK+qBMmLECG3cuFG7du3SmjVr5Pf7VVxcrNOnT8vv90uSsrKywr4mKysrdFtHqqur5fV6Q5fc3NxoTxsAABgS9UApLS3VQw89pMLCQo0bN047duyQJG3YsCE0xuPxhH2Nc67dtq9asmSJAoFA6NLQ0BDtaQMAAENi/jbj1NRUFRYW6ujRo6F381x+tqSpqandWZWvSklJUXp6etgFAAD0XDEPlGAwqCNHjig7O1v5+fny+XyqqakJ3d7W1qba2loVFxfHeioAAKCbSIz2HS5atEj333+/rr/+ejU1NelXv/qVmpubNXPmTHk8HpWVlamqqkpDhw7V0KFDVVVVpX79+mnGjBnRngoAAOimoh4oJ0+e1I9+9COdOnVKAwcO1MiRI7Vv3z7l5eVJkhYvXqxz587piSee0JkzZzRixAjt3r1baWlp0Z4KAADopjzOORfvSUSqublZXq9XgUAg7q9HGVKxI677t+TEsonxngIAwLBIfn7zt3gAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDoECAADMSYz3BIBur9Ib7xl0T5WBeM8AgGEECqJmSMWOqx57YtnEGM4EANDd8SseAABgDoECAADMIVAAAIA5BAoAADCHQAEAAOYQKAAAwBwCBQAAmEOgAAAAcwgUAABgDp8kCwCXs/znC/gTAeglOIMCAADMIVAAAIA5BAoAADCH16AAQHdi9fUxvDYGUcYZFAAAYA6BAgAAzCFQAACAObwGpQNDKnbEewoAAPRqnEEBAADmECgAAMAcAgUAAJjDa1AQF5G8zufEsokxnAkAwCICBd2H1Q+oAmD7+ORD5LolfsUDAADMIVAAAIA5BAoAADCHQAEAAObENVCef/555efnq0+fPioqKtK7774bz+nAqCEVO/h0XwDoZeIWKNu2bVNZWZmeeuopHThwQHfddZdKS0v16aefxmtKAADAiLi9zXjFihWaPXu2Hn30UUnSs88+q127dmn16tWqrq6O17R6vRN9ZsR7CgAAxCdQ2traVF9fr4qKirDtJSUlqqurazc+GAwqGAyGrgcC/3tPe3Nzc0zm92Xwi5jcb3fQ7HHxngJ6ixgdv1ER5DjoUSx/r1UPjvcMrmzJyajf5aWf28598zEWl0A5deqULl68qKysrLDtWVlZ8vv97cZXV1fr6aefbrc9Nzc3ZnPsrQx/1BJ6mmV8t6GL8L3WOTF83FpaWuT1fv39x/WTZD0eT9h151y7bZK0ZMkSlZeXh65/+eWX+s9//qMBAwZ0OL4jzc3Nys3NVUNDg9LT069t4t1Eb1tzb1uvxJp7w5p723ql3rfm3rRe55xaWlqUk5PzjWPjEiiZmZlKSEhod7akqamp3VkVSUpJSVFKSkrYtm9961ud2nd6enqP/wa4XG9bc29br8Sae4Petl6p9625t6z3m86cXBKXd/EkJyerqKhINTU1YdtrampUXFwcjykBAABD4vYrnvLycj388MMaPny4Ro0apT/+8Y/69NNP9fjjj8drSgAAwIi4Bcq0adN0+vRp/eIXv1BjY6MKCgq0c+dO5eXlxWR/KSkpWrp0abtfFfVkvW3NvW29EmvuDXrbeqXet+bett6r5XFX814fAACALsTf4gEAAOYQKAAAwBwCBQAAmEOgAAAAc7pNoDz//PPKz89Xnz59VFRUpHffffdrx9fW1qqoqEh9+vTRd77zHf3hD38Iu339+vXyeDztLv/973+vab/RFO01jxkzpsM1T5w4MTSmsrKy3e0+ny8m67tcJOttbGzUjBkzdOONN+q6665TWVlZh+Nefvll3XTTTUpJSdFNN92k7du3X9N+oy3aa16zZo3uuusu9e/fX/3799e4ceP03nvvhY2J53MsRX/N1o/laK/X+nEsRbbmV155RePHj9fAgQOVnp6uUaNGadeuXe3GWT6Wo73e7nAcdwnXDWzdutUlJSW5NWvWuMOHD7sFCxa41NRU969//avD8ceOHXP9+vVzCxYscIcPH3Zr1qxxSUlJ7qWXXgqNWbdunUtPT3eNjY1hl2vZbzTFYs2nT58OW+vBgwddQkKCW7duXWjM0qVL3c033xw2rqmpKdbLjXi9x48fdz/5yU/chg0b3Pe//323YMGCdmPq6upcQkKCq6qqckeOHHFVVVUuMTHR7du3r9P7jaZYrHnGjBnuueeecwcOHHBHjhxxjzzyiPN6ve7kyZOhMfF6jp2LzZotH8uxWK/l49i5yNe8YMEC98wzz7j33nvP/eMf/3BLlixxSUlJ7u9//3tojOVjORbrtX4cd5VuESi33367e/zxx8O2DRs2zFVUVHQ4fvHixW7YsGFh2+bMmeNGjhwZur5u3Trn9Xqjut9oisWaL7dy5UqXlpbmWltbQ9uWLl3qbrnlls5PvJOu5bEePXp0h/8hnzp1qrvvvvvCtk2YMMFNnz49Kvu9VrFY8+UuXLjg0tLS3IYNG0Lb4vUcOxebNVs+lrviObZ0HDsXncf6pptuck8//XTouuVjORbrvZy147irmP8VT1tbm+rr61VSUhK2vaSkRHV1dR1+zV//+td24ydMmKD3339f58+fD21rbW1VXl6eBg8erEmTJunAgQPXtN9oieWav2rt2rWaPn26UlNTw7YfPXpUOTk5ys/P1/Tp03Xs2LFrWM03i9VjfaXH5NJ9drfnuDO++OILnT9/XhkZGWHbu/o5lmK7ZovHclft18pxLEVnzV9++aVaWlrCvmetHsuxWu/lLB3HXcl8oJw6dUoXL15s90cEs7Ky2v2xwUv8fn+H4y9cuKBTp05JkoYNG6b169fr9ddf15YtW9SnTx/dcccdOnr0aKf3Gy2xWvNXvffeezp48KAeffTRsO0jRozQxo0btWvXLq1Zs0Z+v1/FxcU6ffr0Na7qymL1WF/pMbl0n93tOe6MiooKffvb39a4ceNC2+LxHEuxW7PVY7kr9mvpOJais+bf/va3Onv2rKZOnRraZvVYjtV6L2fpOO5Kcfuo+0h5PJ6w6865dtu+afxXt48cOVIjR44M3X7HHXfo1ltv1e9//3v97ne/6/R+oynaa/6qtWvXqqCgQLfffnvY9tLS0tC/CwsLNWrUKH33u9/Vhg0bVF5eHvEaIhGLx/pq7rM7PceRWL58ubZs2aI9e/aoT58+oe3xfI6l6K/Z+rEcy/1aPI6lzq95y5Ytqqys1GuvvaZBgwZFfJ/d7Tn+uvVeYvU47grmz6BkZmYqISGhXY02NTW1q9ZLfD5fh+MTExM1YMCADr/muuuu02233Rb6v67O7DdaYr3mL774Qlu3bm33f10dSU1NVWFhYehxiYVYPdZXekwu3Wd3e44j8Zvf/EZVVVXavXu3vve9733t2K54jqWue7ytHMux3q+141i6tjVv27ZNs2fP1osvvhh2pkCyeyzHar2XWDyOu5L5QElOTlZRUZFqamrCttfU1Ki4uLjDrxk1alS78bt379bw4cOVlJTU4dc45/TBBx8oOzu70/uNlliv+cUXX1QwGNSPf/zjb5xLMBjUkSNHQo9LLMTqsb7SY3LpPrvbc3y1fv3rX+uXv/yl3njjDQ0fPvwbx3fFcyx13eNt5ViO9X6tHcdS59e8ZcsWzZo1S5s3bw57u/QlVo/lWK1Xsnscd6muflVuZ1x6G9fatWvd4cOHXVlZmUtNTXUnTpxwzjlXUVHhHn744dD4S2+5/elPf+oOHz7s1q5d2+4tt5WVle6NN95w//znP92BAwfcI4884hITE93f/va3q95vd1vzJXfeeaebNm1ah/tduHCh27Nnjzt27Jjbt2+fmzRpkktLS4v5miNdr3POHThwwB04cMAVFRW5GTNmuAMHDrhDhw6Fbv/LX/7iEhIS3LJly9yRI0fcsmXLrvjWxO7wHF/Nmp955hmXnJzsXnrppbC3H7a0tITGxOs5jtWaLR/LsVjvJRaPY+ciX/PmzZtdYmKie+6558K+Zz///PPQGMvHcizWa/047irdIlCcc+65555zeXl5Ljk52d16662utrY2dNvMmTPd6NGjw8bv2bPH/eAHP3DJycluyJAhbvXq1WG3l5WVueuvv94lJye7gQMHupKSEldXVxfRfmMt2mt2zrmPP/7YSXK7d+/ucJ/Tpk1z2dnZLikpyeXk5LgpU6Z0+B/HWIh0vZLaXfLy8sLG/PnPf3Y33nijS0pKcsOGDXMvv/xyRPuNtWivOS8vr8MxS5cuDY2J53PsXPTXbP1YjsX3teXj2LnI1jx69OgO1zxz5syw+7R8LEd7vd3hOO4KHuf+3yspAQAAjDD/GhQAAND7ECgAAMAcAgUAAJhDoAAAAHMIFAAAYA6BAgAAzCFQAACAOQQKAAAwh0ABAADmECgAAMAcAgUAAJhDoAAAAHP+D8EdnCHboeFmAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.hist(monitor.Phi)\n",
    "plt.hist(max_dis_rates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Investigate behavior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.8500, device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "from bayesian_dpddm.monitors.utils import sample_from_dataset, temperature_scaling\n",
    "import torch.nn.functional as F\n",
    "X, y = sample_from_dataset(dataset=dataset['dpddm_train'], n_samples=100, replace=True)\n",
    "y_pseudo = monitor.get_pseudolabels(X)\n",
    "print(torch.sum(y.to(device) == y_pseudo)/len(y_pseudo))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get max disagreement rates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "X, y_pseudo = X.to(device), y_pseudo.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = monitor.model\n",
    "with torch.no_grad():\n",
    "    output = model.get_features(X)\n",
    "    ll_dist = model.out_layer.logit_predictive(output)\n",
    "    ll_dist.loc = ll_dist.loc\n",
    "    ll_dist.scale = ll_dist.scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [],
   "source": [
    "logits_samples = ll_dist.rsample(sample_shape=torch.Size([5000]))\n",
    "logits_samples = temperature_scaling(logits_samples, temp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [],
   "source": [
    "dist = torch.distributions.Categorical(logits=logits_samples)\n",
    "y_hat = dist.sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.2800000011920929"
      ]
     },
     "execution_count": 109,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tmp = []\n",
    "for i in range(5000):\n",
    "    rate = (y_hat[i] != y_pseudo).float().mean().item()\n",
    "    tmp.append(rate)\n",
    "max(tmp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5000, 100])"
      ]
     },
     "execution_count": 110,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_tile = torch.tile(y_pseudo, (5000, 1)).cuda()\n",
    "y_tile.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5000, 100])"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_hat.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.2800000011920929"
      ]
     },
     "execution_count": 112,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dis_mat = (y_hat != y_tile)\n",
    "dis_rate = dis_mat.sum(dim=-1)/len(y)\n",
    "torch.max(dis_rate).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5000])"
      ]
     },
     "execution_count": 113,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dis_rate.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.1700, 0.1600, 0.1600,  ..., 0.1900, 0.1500, 0.1800], device='cuda:0')"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dis_rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Normal(loc: torch.Size([100, 10]), scale: torch.Size([100, 10]))"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ll_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0456, 1.1385, 0.8966, 0.8294, 0.6944, 0.6931, 1.0488, 0.5736, 1.1509,\n",
       "         1.0840],\n",
       "        [0.6187, 0.8241, 0.9640, 1.0311, 1.0410, 1.0718, 1.0056, 1.0603, 0.4203,\n",
       "         0.7717],\n",
       "        [2.0480, 2.3146, 1.5980, 0.5035, 1.4490, 0.8518, 1.3738, 1.6512, 2.1593,\n",
       "         2.1676],\n",
       "        [0.7446, 0.4925, 0.9979, 0.9567, 0.9915, 0.9522, 0.9558, 0.8848, 0.7108,\n",
       "         0.3895],\n",
       "        [4.1144, 0.5396, 4.5568, 4.6295, 4.5940, 4.5411, 4.3329, 4.5622, 3.2926,\n",
       "         1.3016],\n",
       "        [1.2600, 1.2357, 0.8931, 0.7486, 0.9267, 0.9764, 0.6367, 1.2011, 1.2369,\n",
       "         1.2705],\n",
       "        [3.0622, 4.1787, 5.4772, 5.6666, 5.8027, 5.8814, 5.4171, 5.8891, 0.6446,\n",
       "         3.9505],\n",
       "        [2.9098, 2.5881, 1.6841, 1.7148, 1.9675, 2.4209, 0.4680, 2.8712, 2.7352,\n",
       "         2.8540],\n",
       "        [0.9475, 1.5933, 0.4696, 1.0895, 0.8584, 1.0915, 0.8754, 1.2814, 1.5120,\n",
       "         1.6446],\n",
       "        [0.8730, 0.9582, 0.7635, 0.6031, 0.6203, 0.5729, 0.7834, 0.6145, 0.9320,\n",
       "         0.9017],\n",
       "        [0.7297, 1.6583, 0.8762, 1.5823, 1.4265, 1.6907, 1.6331, 1.6108, 1.1934,\n",
       "         1.4808],\n",
       "        [2.2011, 2.6673, 1.4110, 1.6008, 0.4396, 1.7964, 1.7248, 1.0498, 2.5230,\n",
       "         2.6539],\n",
       "        [1.4115, 1.9872, 2.4990, 2.6370, 2.6841, 2.7446, 2.5418, 2.7310, 0.5077,\n",
       "         1.8562],\n",
       "        [0.7002, 2.6522, 1.4595, 2.7696, 2.5171, 3.0128, 2.9448, 2.7424, 1.5076,\n",
       "         2.1785],\n",
       "        [0.4850, 1.9213, 1.1395, 2.0744, 1.9041, 2.2653, 2.2296, 2.0348, 0.9894,\n",
       "         1.5263],\n",
       "        [0.7011, 1.7350, 1.3727, 1.9781, 1.8769, 2.1296, 2.0620, 1.9727, 0.8184,\n",
       "         1.4361],\n",
       "        [3.2909, 3.5140, 2.2947, 0.7351, 2.2629, 1.0142, 2.3558, 2.2970, 3.4022,\n",
       "         3.3593],\n",
       "        [0.4778, 1.8703, 1.1461, 2.0363, 1.8743, 2.2208, 2.1814, 1.9996, 0.9418,\n",
       "         1.4881],\n",
       "        [2.4132, 2.3905, 1.4444, 0.6421, 1.6232, 0.4667, 1.8715, 1.3472, 2.4370,\n",
       "         2.3574],\n",
       "        [1.2326, 0.6956, 1.7255, 1.6463, 1.7108, 1.6365, 1.6428, 1.5040, 1.1643,\n",
       "         0.4371],\n",
       "        [1.6189, 0.8378, 2.3064, 2.1965, 2.2857, 2.1823, 2.1911, 1.9983, 1.5218,\n",
       "         0.4040],\n",
       "        [0.9071, 0.9689, 0.8543, 0.7927, 0.6285, 0.6711, 0.9628, 0.4886, 0.9736,\n",
       "         0.9065],\n",
       "        [3.5391, 3.5042, 2.0932, 0.8548, 2.3613, 0.5530, 2.7315, 1.9462, 3.5739,\n",
       "         3.4557],\n",
       "        [2.2741, 1.1650, 3.2454, 3.0902, 3.2160, 3.0697, 3.0821, 2.8099, 2.1363,\n",
       "         0.5395],\n",
       "        [4.0548, 4.0143, 2.3865, 0.9368, 2.6965, 0.5620, 3.1238, 2.2160, 4.0947,\n",
       "         3.9586],\n",
       "        [3.4934, 3.1043, 2.0069, 2.0442, 2.3511, 2.9015, 0.4761, 3.4462, 3.2824,\n",
       "         3.4263],\n",
       "        [1.0180, 1.2049, 0.7173, 0.7863, 0.4340, 0.8610, 0.8344, 0.5955, 1.1460,\n",
       "         1.1983],\n",
       "        [0.4533, 1.7441, 1.0540, 1.8907, 1.7378, 2.0628, 2.0280, 1.8556, 0.8918,\n",
       "         1.3885],\n",
       "        [0.4006, 0.4234, 0.3969, 0.4118, 0.4092, 0.4188, 0.4190, 0.4155, 0.4056,\n",
       "         0.4142],\n",
       "        [0.5667, 0.5860, 0.4745, 0.3987, 0.4732, 0.4325, 0.4550, 0.5092, 0.5737,\n",
       "         0.5750],\n",
       "        [1.8952, 2.5679, 3.3529, 3.4676, 3.5502, 3.5981, 3.3166, 3.6029, 0.5171,\n",
       "         2.4299],\n",
       "        [0.4164, 1.2769, 0.7877, 1.3748, 1.2659, 1.4966, 1.4736, 1.3493, 0.6981,\n",
       "         1.0275],\n",
       "        [0.7446, 0.4925, 0.9979, 0.9567, 0.9915, 0.9522, 0.9558, 0.8848, 0.7108,\n",
       "         0.3895],\n",
       "        [1.2364, 1.2484, 0.8708, 0.6256, 0.8402, 0.5207, 1.0591, 0.6912, 1.2708,\n",
       "         1.2143],\n",
       "        [1.0645, 1.5610, 1.8725, 2.0255, 2.0469, 2.1154, 1.9709, 2.0891, 0.4982,\n",
       "         1.4377],\n",
       "        [0.4859, 1.7729, 1.0651, 1.9159, 1.7596, 2.0899, 2.0568, 1.8785, 0.9279,\n",
       "         1.4132],\n",
       "        [0.8092, 0.9637, 0.5440, 0.5855, 0.6513, 0.5685, 0.7033, 0.7141, 0.9483,\n",
       "         0.9698],\n",
       "        [0.9141, 1.2123, 1.4623, 1.5326, 1.5565, 1.5862, 1.4704, 1.5918, 0.5790,\n",
       "         1.1582],\n",
       "        [2.1408, 1.9092, 1.2579, 1.2801, 1.4607, 1.7875, 0.4374, 2.1127, 2.0151,\n",
       "         2.1010],\n",
       "        [2.5819, 2.5898, 1.5572, 1.6518, 1.4825, 2.1388, 1.2064, 2.2190, 2.6077,\n",
       "         2.7292],\n",
       "        [0.5485, 0.5078, 0.5042, 0.5007, 0.5185, 0.5406, 0.4648, 0.5618, 0.5350,\n",
       "         0.5173],\n",
       "        [1.7677, 1.8920, 1.5840, 1.4115, 0.9859, 1.0746, 1.8508, 0.4773, 1.9339,\n",
       "         1.7511],\n",
       "        [1.5469, 1.6288, 1.3106, 1.1291, 0.9277, 0.8759, 1.5414, 0.5963, 1.6628,\n",
       "         1.5288],\n",
       "        [2.4931, 2.4692, 1.4905, 0.6586, 1.6757, 0.4759, 1.9323, 1.3903, 2.5175,\n",
       "         2.4352],\n",
       "        [0.6785, 1.1338, 0.3988, 0.7975, 0.6534, 0.7826, 0.6761, 0.9171, 1.0760,\n",
       "         1.1651],\n",
       "        [1.6863, 3.1751, 0.6154, 2.1486, 1.6414, 2.0985, 1.7018, 2.5219, 2.9939,\n",
       "         3.2769],\n",
       "        [1.0369, 0.9373, 0.6648, 0.6727, 0.7468, 0.8842, 0.3893, 1.0237, 0.9819,\n",
       "         1.0186],\n",
       "        [1.9688, 2.6610, 3.4721, 3.5910, 3.6775, 3.7286, 3.4320, 3.7354, 0.5498,\n",
       "         2.5197],\n",
       "        [1.4766, 1.6949, 1.2024, 0.4408, 1.0712, 0.6964, 0.9832, 1.2522, 1.5649,\n",
       "         1.5789],\n",
       "        [1.7601, 1.8940, 1.2711, 0.5024, 1.2379, 0.6472, 1.2693, 1.2797, 1.8238,\n",
       "         1.8060],\n",
       "        [3.1686, 3.8832, 2.0493, 2.3803, 0.8180, 2.6691, 2.5657, 1.6138, 3.6431,\n",
       "         3.8483],\n",
       "        [0.6882, 0.7707, 0.6490, 0.6509, 0.5339, 0.5907, 0.7676, 0.4670, 0.7489,\n",
       "         0.7155],\n",
       "        [1.8792, 1.9630, 1.2706, 0.5098, 1.2968, 0.5790, 1.3912, 1.2533, 1.9286,\n",
       "         1.8937],\n",
       "        [1.5539, 0.8115, 2.1941, 2.0935, 2.1757, 2.0799, 2.0856, 1.9106, 1.4547,\n",
       "         0.4356],\n",
       "        [1.3462, 1.8163, 2.3612, 2.4427, 2.5001, 2.5343, 2.3379, 2.5374, 0.4241,\n",
       "         1.7191],\n",
       "        [0.7344, 1.2445, 0.4307, 0.8803, 0.7148, 0.8639, 0.7374, 1.0100, 1.1799,\n",
       "         1.2798],\n",
       "        [1.4537, 1.5693, 1.0563, 0.4967, 1.0183, 0.6050, 1.0610, 1.0591, 1.5108,\n",
       "         1.5004],\n",
       "        [1.4618, 1.5842, 1.3030, 1.1864, 0.7978, 0.9441, 1.5179, 0.4450, 1.6075,\n",
       "         1.4753],\n",
       "        [2.2613, 1.2129, 3.2568, 3.1093, 3.2347, 3.0941, 3.0963, 2.8422, 2.1007,\n",
       "         0.6363],\n",
       "        [0.7179, 1.1831, 0.4725, 0.8625, 0.7145, 0.8510, 0.7353, 0.9768, 1.1194,\n",
       "         1.2127],\n",
       "        [1.4846, 2.0053, 2.6124, 2.7016, 2.7657, 2.8030, 2.5848, 2.8068, 0.4433,\n",
       "         1.8981],\n",
       "        [0.8625, 0.9410, 0.7916, 0.7575, 0.6062, 0.6518, 0.8988, 0.5163, 0.9498,\n",
       "         0.8923],\n",
       "        [2.5041, 2.2303, 1.4570, 1.4836, 1.6976, 2.0859, 0.4473, 2.4701, 2.3556,\n",
       "         2.4573],\n",
       "        [0.4631, 0.5146, 0.4312, 0.4587, 0.4755, 0.4716, 0.5039, 0.4770, 0.4771,\n",
       "         0.4920],\n",
       "        [0.4740, 0.4840, 0.4655, 0.4573, 0.4388, 0.4447, 0.4842, 0.4253, 0.4866,\n",
       "         0.4709],\n",
       "        [0.6722, 0.6803, 0.5306, 0.4794, 0.5318, 0.5574, 0.4608, 0.6341, 0.6732,\n",
       "         0.6874],\n",
       "        [4.0874, 0.5080, 4.5278, 4.6000, 4.5647, 4.5119, 4.3050, 4.5328, 3.2696,\n",
       "         1.2823],\n",
       "        [0.8730, 0.9582, 0.7635, 0.6031, 0.6203, 0.5729, 0.7834, 0.6145, 0.9320,\n",
       "         0.9017],\n",
       "        [3.2667, 1.0307, 3.5092, 3.5606, 3.5521, 3.5424, 3.2990, 3.5898, 2.6713,\n",
       "         1.4087],\n",
       "        [1.6661, 2.3243, 2.9580, 3.1032, 3.1642, 3.2272, 2.9842, 3.2171, 0.5399,\n",
       "         2.1785],\n",
       "        [2.1145, 1.0833, 3.0176, 2.8733, 2.9903, 2.8543, 2.8659, 2.6128, 1.9865,\n",
       "         0.5013],\n",
       "        [0.5914, 0.8289, 0.4554, 0.6383, 0.5673, 0.6419, 0.5644, 0.7168, 0.7967,\n",
       "         0.8489],\n",
       "        [1.1436, 1.5078, 1.9206, 2.0545, 2.0796, 2.1338, 2.0131, 2.0817, 0.7365,\n",
       "         1.3610],\n",
       "        [0.4450, 1.0289, 0.6866, 1.1029, 1.0245, 1.1932, 1.1750, 1.0857, 0.6142,\n",
       "         0.8484],\n",
       "        [1.7861, 0.9342, 2.4961, 2.3851, 2.4758, 2.3686, 2.3732, 2.1813, 1.6683,\n",
       "         0.5390],\n",
       "        [0.4755, 0.4958, 0.4593, 0.4370, 0.4571, 0.4452, 0.4549, 0.4669, 0.4835,\n",
       "         0.4859],\n",
       "        [1.6957, 1.9195, 1.0856, 1.0965, 0.7642, 1.2346, 1.2805, 0.9857, 1.8557,\n",
       "         1.9182],\n",
       "        [0.9882, 0.9496, 1.0677, 0.8960, 1.0130, 0.9144, 0.9998, 0.9727, 1.0064,\n",
       "         0.8601],\n",
       "        [1.8941, 3.6277, 0.5728, 2.4359, 1.8411, 2.3771, 1.9124, 2.8704, 3.4180,\n",
       "         3.7457],\n",
       "        [2.4376, 2.5648, 1.9639, 1.5590, 1.4916, 1.2182, 2.2895, 1.0993, 2.5955,\n",
       "         2.4205],\n",
       "        [0.8812, 1.5901, 0.4201, 1.0957, 0.8585, 1.0727, 0.8879, 1.2745, 1.5021,\n",
       "         1.6385],\n",
       "        [0.4162, 0.4219, 0.4049, 0.4054, 0.4075, 0.4114, 0.4070, 0.4154, 0.4187,\n",
       "         0.4192],\n",
       "        [1.8791, 2.5546, 3.3407, 3.4557, 3.5383, 3.5863, 3.3047, 3.5911, 0.4645,\n",
       "         2.4162],\n",
       "        [1.8300, 2.1761, 1.1801, 1.3253, 0.5340, 1.5020, 1.3945, 0.9887, 2.0694,\n",
       "         2.1745],\n",
       "        [2.4346, 2.4412, 1.4962, 0.6330, 1.6431, 0.5241, 1.8662, 1.4168, 2.4675,\n",
       "         2.3949],\n",
       "        [1.9004, 2.2936, 1.2414, 1.3990, 0.4901, 1.5619, 1.5020, 0.9488, 2.1717,\n",
       "         2.2822],\n",
       "        [1.8555, 1.0084, 2.6501, 2.5413, 2.6390, 2.5296, 2.5387, 2.3150, 1.7500,\n",
       "         0.5203],\n",
       "        [0.9162, 0.8365, 0.6075, 0.6147, 0.6718, 0.7882, 0.3984, 0.9016, 0.8721,\n",
       "         0.9030],\n",
       "        [1.1047, 1.3060, 1.7371, 1.7957, 1.8291, 1.8508, 1.7166, 1.8560, 0.6303,\n",
       "         1.2521],\n",
       "        [0.4716, 1.9165, 1.1357, 2.0739, 1.9023, 2.2645, 2.2283, 2.0327, 0.9818,\n",
       "         1.5208],\n",
       "        [1.8060, 1.9963, 1.0337, 1.3355, 1.2763, 1.6267, 0.8630, 1.8783, 2.0000,\n",
       "         2.1226],\n",
       "        [1.2209, 1.3090, 1.0977, 1.0009, 0.7147, 0.8098, 1.2634, 0.4684, 1.3318,\n",
       "         1.2216],\n",
       "        [2.3444, 3.1974, 4.1890, 4.3338, 4.4378, 4.4980, 4.1433, 4.5040, 0.5093,\n",
       "         3.0229],\n",
       "        [2.1758, 1.9422, 1.2875, 1.3097, 1.4910, 1.8197, 0.4808, 2.1474, 2.0489,\n",
       "         2.1356],\n",
       "        [4.2597, 3.7829, 2.4394, 2.4851, 2.8614, 3.5354, 0.5275, 4.2020, 4.0013,\n",
       "         4.1772],\n",
       "        [1.4353, 1.4224, 0.8869, 0.4682, 0.9882, 0.3921, 1.1278, 0.8357, 1.4489,\n",
       "         1.4031],\n",
       "        [0.4266, 1.6216, 0.9694, 1.7537, 1.6104, 1.9137, 1.8834, 1.7197, 0.8412,\n",
       "         1.2899],\n",
       "        [2.8240, 2.8473, 1.7539, 0.7168, 1.9070, 0.6226, 2.1516, 1.6702, 2.8669,\n",
       "         2.7867],\n",
       "        [2.1899, 2.5229, 1.7720, 0.5106, 1.5692, 0.9602, 1.4223, 1.8495, 2.3236,\n",
       "         2.3444],\n",
       "        [2.5532, 3.4850, 4.5680, 4.7261, 4.8396, 4.9053, 4.5181, 4.9118, 0.5326,\n",
       "         3.2945]], device='cuda:0')"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ll_dist.scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.0934,  1.2962,  0.8039,  0.6878,  0.4823,  0.4804,  1.1000,  0.3290,\n",
       "          1.3247,  1.1751],\n",
       "        [ 0.3828,  0.6791,  0.9293,  1.0632,  1.0837,  1.1488,  1.0113,  1.1243,\n",
       "          0.1766,  0.5955],\n",
       "        [ 4.1941,  5.3574,  2.5535,  0.2535,  2.0995,  0.7256,  1.8873,  2.7264,\n",
       "          4.6627,  4.6985],\n",
       "        [ 0.5544,  0.2426,  0.9957,  0.9153,  0.9831,  0.9067,  0.9135,  0.7829,\n",
       "          0.5053,  0.1517],\n",
       "        [16.9282,  0.2911, 20.7648, 21.4322, 21.1051, 20.6214, 18.7744, 20.8134,\n",
       "         10.8414,  1.6941],\n",
       "        [ 1.5877,  1.5271,  0.7976,  0.5604,  0.8587,  0.9534,  0.4054,  1.4425,\n",
       "          1.5300,  1.6142],\n",
       "        [ 9.3772, 17.4614, 29.9999, 32.1098, 33.6710, 34.5903, 29.3450, 34.6817,\n",
       "          0.4156, 15.6064],\n",
       "        [ 8.4669,  6.6981,  2.8361,  2.9405,  3.8711,  5.8608,  0.2190,  8.2439,\n",
       "          7.4816,  8.1455],\n",
       "        [ 0.8978,  2.5386,  0.2205,  1.1870,  0.7369,  1.1913,  0.7664,  1.6420,\n",
       "          2.2860,  2.7047],\n",
       "        [ 0.7622,  0.9181,  0.5830,  0.3637,  0.3848,  0.3282,  0.6137,  0.3776,\n",
       "          0.8687,  0.8130],\n",
       "        [ 0.5324,  2.7500,  0.7678,  2.5038,  2.0349,  2.8583,  2.6671,  2.5947,\n",
       "          1.4243,  2.1927],\n",
       "        [ 4.8449,  7.1144,  1.9910,  2.5626,  0.1933,  3.2271,  2.9749,  1.1020,\n",
       "          6.3654,  7.0434],\n",
       "        [ 1.9924,  3.9490,  6.2452,  6.9536,  7.2043,  7.5327,  6.4608,  7.4584,\n",
       "          0.2578,  3.4457],\n",
       "        [ 0.4902,  7.0342,  2.1302,  7.6709,  6.3356,  9.0769,  8.6718,  7.5209,\n",
       "          2.2729,  4.7460],\n",
       "        [ 0.2352,  3.6914,  1.2984,  4.3032,  3.6255,  5.1314,  4.9712,  4.1403,\n",
       "          0.9788,  2.3296],\n",
       "        [ 0.4915,  3.0102,  1.8843,  3.9131,  3.5227,  4.5353,  4.2517,  3.8916,\n",
       "          0.6697,  2.0624],\n",
       "        [10.8298, 12.3480,  5.2655,  0.5404,  5.1208,  1.0286,  5.5499,  5.2763,\n",
       "         11.5747, 11.2847],\n",
       "        [ 0.2283,  3.4979,  1.3135,  4.1463,  3.5128,  4.9321,  4.7583,  3.9985,\n",
       "          0.8871,  2.2145],\n",
       "        [ 5.8235,  5.7143,  2.0864,  0.4123,  2.6349,  0.2178,  3.5025,  1.8150,\n",
       "          5.9388,  5.5573],\n",
       "        [ 1.5192,  0.4839,  2.9773,  2.7105,  2.9268,  2.6780,  2.6988,  2.2621,\n",
       "          1.3555,  0.1911],\n",
       "        [ 2.6208,  0.7018,  5.3193,  4.8245,  5.2242,  4.7622,  4.8010,  3.9934,\n",
       "          2.3160,  0.1632],\n",
       "        [ 0.8228,  0.9388,  0.7299,  0.6283,  0.3951,  0.4504,  0.9270,  0.2388,\n",
       "          0.9479,  0.8217],\n",
       "        [12.5256, 12.2792,  4.3814,  0.7308,  5.5758,  0.3058,  7.4610,  3.7878,\n",
       "         12.7728, 11.9418],\n",
       "        [ 5.1715,  1.3571, 10.5326,  9.5490, 10.3429,  9.4233,  9.4993,  7.8956,\n",
       "          4.5639,  0.2910],\n",
       "        [16.4413, 16.1144,  5.6955,  0.8777,  7.2708,  0.3158,  9.7584,  4.9107,\n",
       "         16.7668, 15.6701],\n",
       "        [12.2040,  9.6369,  4.0275,  4.1789,  5.5277,  8.4185,  0.2267, 11.8765,\n",
       "         10.7742, 11.7397],\n",
       "        [ 1.0364,  1.4517,  0.5145,  0.6183,  0.1883,  0.7414,  0.6962,  0.3546,\n",
       "          1.3133,  1.4360],\n",
       "        [ 0.2055,  3.0419,  1.1110,  3.5748,  3.0199,  4.2551,  4.1127,  3.4432,\n",
       "          0.7953,  1.9279],\n",
       "        [ 0.1605,  0.1792,  0.1575,  0.1695,  0.1674,  0.1754,  0.1755,  0.1726,\n",
       "          0.1645,  0.1715],\n",
       "        [ 0.3211,  0.3434,  0.2252,  0.1589,  0.2239,  0.1870,  0.2071,  0.2592,\n",
       "          0.3291,  0.3306],\n",
       "        [ 3.5919,  6.5940, 11.2418, 12.0243, 12.6042, 12.9461, 10.9997, 12.9808,\n",
       "          0.2674,  5.9043],\n",
       "        [ 0.1734,  1.6305,  0.6204,  1.8900,  1.6026,  2.2398,  2.1715,  1.8206,\n",
       "          0.4874,  1.0557],\n",
       "        [ 0.5544,  0.2426,  0.9957,  0.9153,  0.9831,  0.9067,  0.9135,  0.7829,\n",
       "          0.5053,  0.1517],\n",
       "        [ 1.5288,  1.5584,  0.7582,  0.3913,  0.7060,  0.2711,  1.1217,  0.4778,\n",
       "          1.6150,  1.4745],\n",
       "        [ 1.1331,  2.4366,  3.5062,  4.1025,  4.1898,  4.4749,  3.8843,  4.3645,\n",
       "          0.2482,  2.0669],\n",
       "        [ 0.2361,  3.1432,  1.1344,  3.6706,  3.0963,  4.3675,  4.2304,  3.5287,\n",
       "          0.8609,  1.9973],\n",
       "        [ 0.6548,  0.9287,  0.2960,  0.3428,  0.4242,  0.3232,  0.4946,  0.5099,\n",
       "          0.8993,  0.9404],\n",
       "        [ 0.8355,  1.4696,  2.1384,  2.3488,  2.4228,  2.5162,  2.1620,  2.5337,\n",
       "          0.3352,  1.3415],\n",
       "        [ 4.5832,  3.6450,  1.5824,  1.6388,  2.1336,  3.1950,  0.1914,  4.4634,\n",
       "          4.0604,  4.4142],\n",
       "        [ 6.6661,  6.7072,  2.4248,  2.7283,  2.1977,  4.5747,  1.4554,  4.9241,\n",
       "          6.8000,  7.4483],\n",
       "        [ 0.3009,  0.2578,  0.2543,  0.2507,  0.2688,  0.2923,  0.2161,  0.3156,\n",
       "          0.2862,  0.2676],\n",
       "        [ 3.1249,  3.5796,  2.5090,  1.9922,  0.9720,  1.1548,  3.4254,  0.2278,\n",
       "          3.7398,  3.0665],\n",
       "        [ 2.3928,  2.6529,  1.7176,  1.2748,  0.8605,  0.7672,  2.3758,  0.3556,\n",
       "          2.7648,  2.3372],\n",
       "        [ 6.2154,  6.0972,  2.2215,  0.4337,  2.8080,  0.2265,  3.7337,  1.9328,\n",
       "          6.3378,  5.9304],\n",
       "        [ 0.4603,  1.2854,  0.1591,  0.6360,  0.4269,  0.6124,  0.4571,  0.8411,\n",
       "          1.1578,  1.3573],\n",
       "        [ 2.8434, 10.0815,  0.3787,  4.6167,  2.6942,  4.4035,  2.8962,  6.3599,\n",
       "          8.9632, 10.7381],\n",
       "        [ 1.0751,  0.8786,  0.4419,  0.4526,  0.5578,  0.7818,  0.1516,  1.0480,\n",
       "          0.9640,  1.0375],\n",
       "        [ 3.8761,  7.0811, 12.0558, 12.8955, 13.5243, 13.9028, 11.7789, 13.9530,\n",
       "          0.3023,  6.3488],\n",
       "        [ 2.1803,  2.8727,  1.4457,  0.1943,  1.1476,  0.4850,  0.9668,  1.5679,\n",
       "          2.4488,  2.4930],\n",
       "        [ 3.0981,  3.5873,  1.6157,  0.2524,  1.5325,  0.4189,  1.6111,  1.6377,\n",
       "          3.3263,  3.2615],\n",
       "        [10.0400, 15.0789,  4.1997,  5.6660,  0.6690,  7.1239,  6.5830,  2.6045,\n",
       "         13.2723, 14.8091],\n",
       "        [ 0.4736,  0.5940,  0.4213,  0.4237,  0.2851,  0.3489,  0.5893,  0.2181,\n",
       "          0.5608,  0.5119],\n",
       "        [ 3.5315,  3.8532,  1.6143,  0.2599,  1.6817,  0.3352,  1.9354,  1.5707,\n",
       "          3.7196,  3.5861],\n",
       "        [ 2.4146,  0.6586,  4.8140,  4.3828,  4.7338,  4.3259,  4.3496,  3.6505,\n",
       "          2.1161,  0.1897],\n",
       "        [ 1.8124,  3.2988,  5.5751,  5.9668,  6.2505,  6.4229,  5.4656,  6.4386,\n",
       "          0.1798,  2.9553],\n",
       "        [ 0.5393,  1.5488,  0.1855,  0.7749,  0.5109,  0.7463,  0.5438,  1.0202,\n",
       "          1.3921,  1.6380],\n",
       "        [ 2.1133,  2.4628,  1.1157,  0.2468,  1.0369,  0.3660,  1.1258,  1.1217,\n",
       "          2.2826,  2.2513],\n",
       "        [ 2.1369,  2.5096,  1.6979,  1.4076,  0.6364,  0.8913,  2.3040,  0.1980,\n",
       "          2.5841,  2.1765],\n",
       "        [ 5.1133,  1.4710, 10.6067,  9.6679, 10.4635,  9.5735,  9.5869,  8.0782,\n",
       "          4.4130,  0.4049],\n",
       "        [ 0.5154,  1.3998,  0.2232,  0.7440,  0.5104,  0.7242,  0.5407,  0.9541,\n",
       "          1.2530,  1.4708],\n",
       "        [ 2.2040,  4.0211,  6.8245,  7.2987,  7.6488,  7.8569,  6.6811,  7.8780,\n",
       "          0.1965,  3.6029],\n",
       "        [ 0.7438,  0.8854,  0.6267,  0.5738,  0.3674,  0.4248,  0.8079,  0.2666,\n",
       "          0.9022,  0.7962],\n",
       "        [ 6.2707,  4.9744,  2.1228,  2.2010,  2.8819,  4.3511,  0.2001,  6.1012,\n",
       "          5.5489,  6.0383],\n",
       "        [ 0.2145,  0.2648,  0.1859,  0.2104,  0.2261,  0.2224,  0.2539,  0.2276,\n",
       "          0.2276,  0.2421],\n",
       "        [ 0.2247,  0.2343,  0.2167,  0.2091,  0.1925,  0.1977,  0.2344,  0.1809,\n",
       "          0.2368,  0.2217],\n",
       "        [ 0.4519,  0.4628,  0.2815,  0.2298,  0.2828,  0.3107,  0.2123,  0.4021,\n",
       "          0.4531,  0.4726],\n",
       "        [16.7067,  0.2581, 20.5006, 21.1596, 20.8364, 20.3577, 18.5326, 20.5466,\n",
       "         10.6903,  1.6443],\n",
       "        [ 0.7622,  0.9181,  0.5830,  0.3637,  0.3848,  0.3282,  0.6137,  0.3776,\n",
       "          0.8687,  0.8130],\n",
       "        [10.6710,  1.0623, 12.3143, 12.6777, 12.6176, 12.5484, 10.8833, 12.8866,\n",
       "          7.1359,  1.9846],\n",
       "        [ 2.7758,  5.4022,  8.7497,  9.6300, 10.0121, 10.4149,  8.9052, 10.3495,\n",
       "          0.2915,  4.7459],\n",
       "        [ 4.4711,  1.1736,  9.1061,  8.2559,  8.9421,  8.1473,  8.2133,  6.8266,\n",
       "          3.9460,  0.2513],\n",
       "        [ 0.3498,  0.6870,  0.2074,  0.4075,  0.3218,  0.4121,  0.3185,  0.5138,\n",
       "          0.6348,  0.7206],\n",
       "        [ 1.3078,  2.2734,  3.6885,  4.2209,  4.3248,  4.5530,  4.0527,  4.3334,\n",
       "          0.5425,  1.8523],\n",
       "        [ 0.1981,  1.0586,  0.4714,  1.2164,  1.0495,  1.4238,  1.3806,  1.1788,\n",
       "          0.3772,  0.7197],\n",
       "        [ 3.1903,  0.8726,  6.2305,  5.6885,  6.1295,  5.6105,  5.6321,  4.7581,\n",
       "          2.7832,  0.2905],\n",
       "        [ 0.2261,  0.2459,  0.2110,  0.1910,  0.2090,  0.1982,  0.2069,  0.2180,\n",
       "          0.2337,  0.2361],\n",
       "        [ 2.8753,  3.6846,  1.1785,  1.2023,  0.5840,  1.5242,  1.6396,  0.9715,\n",
       "          3.4437,  3.6793],\n",
       "        [ 0.9765,  0.9018,  1.1400,  0.8029,  1.0262,  0.8361,  0.9997,  0.9461,\n",
       "          1.0129,  0.7397],\n",
       "        [ 3.5876, 13.1604,  0.3280,  5.9338,  3.3898,  5.6506,  3.6572,  8.2392,\n",
       "         11.6824, 14.0301],\n",
       "        [ 5.9419,  6.5783,  3.8570,  2.4303,  2.2248,  1.4841,  5.2420,  1.2085,\n",
       "          6.7364,  5.8589],\n",
       "        [ 0.7765,  2.5285,  0.1765,  1.2006,  0.7370,  1.1508,  0.7883,  1.6244,\n",
       "          2.2563,  2.6848],\n",
       "        [ 0.1732,  0.1780,  0.1639,  0.1643,  0.1661,  0.1693,  0.1657,  0.1726,\n",
       "          0.1753,  0.1757],\n",
       "        [ 3.5310,  6.5262, 11.1605, 11.9421, 12.5198, 12.8615, 10.9207, 12.8961,\n",
       "          0.2158,  5.8379],\n",
       "        [ 3.3487,  4.7355,  1.3925,  1.7563,  0.2851,  2.2559,  1.9446,  0.9776,\n",
       "          4.2826,  4.7286],\n",
       "        [ 5.9271,  5.9597,  2.2386,  0.4006,  2.6997,  0.2747,  3.4826,  2.0073,\n",
       "          6.0886,  5.7357],\n",
       "        [ 3.6113,  5.2605,  1.5411,  1.9572,  0.2402,  2.4396,  2.2561,  0.9003,\n",
       "          4.7164,  5.2085],\n",
       "        [ 3.4430,  1.0168,  7.0231,  6.4580,  6.9643,  6.3990,  6.4450,  5.3591,\n",
       "          3.0626,  0.2708],\n",
       "        [ 0.8394,  0.6998,  0.3691,  0.3779,  0.4513,  0.6212,  0.1587,  0.8129,\n",
       "          0.7606,  0.8154],\n",
       "        [ 1.2204,  1.7056,  3.0176,  3.2246,  3.3457,  3.4256,  2.9468,  3.4446,\n",
       "          0.3973,  1.5677],\n",
       "        [ 0.2224,  3.6731,  1.2897,  4.3010,  3.6189,  5.1280,  4.9654,  4.1320,\n",
       "          0.9638,  2.3129],\n",
       "        [ 3.2618,  3.9853,  1.0686,  1.7835,  1.6289,  2.6462,  0.7447,  3.5282,\n",
       "          3.9999,  4.5056],\n",
       "        [ 1.4905,  1.7134,  1.2050,  1.0019,  0.5108,  0.6557,  1.5963,  0.2194,\n",
       "          1.7736,  1.4922],\n",
       "        [ 5.4962, 10.2235, 17.5478, 18.7817, 19.6938, 20.2322, 17.1670, 20.2862,\n",
       "          0.2594,  9.1380],\n",
       "        [ 4.7340,  3.7722,  1.6576,  1.7153,  2.2231,  3.3112,  0.2312,  4.6113,\n",
       "          4.1979,  4.5607],\n",
       "        [18.1451, 14.3103,  5.9507,  6.1758,  8.1879, 12.4989,  0.2783, 17.6572,\n",
       "         16.0100, 17.4488],\n",
       "        [ 2.0600,  2.0233,  0.7866,  0.2192,  0.9766,  0.1538,  1.2720,  0.6983,\n",
       "          2.0993,  1.9687],\n",
       "        [ 0.1820,  2.6294,  0.9397,  3.0754,  2.5933,  3.6621,  3.5473,  2.9575,\n",
       "          0.7076,  1.6639],\n",
       "        [ 7.9748,  8.1070,  3.0763,  0.5137,  3.6367,  0.3876,  4.6293,  2.7896,\n",
       "          8.2192,  7.7656],\n",
       "        [ 4.7955,  6.3649,  3.1399,  0.2607,  2.4622,  0.9220,  2.0228,  3.4208,\n",
       "          5.3990,  5.4964],\n",
       "        [ 6.5187, 12.1455, 20.8670, 22.3356, 23.4216, 24.0621, 20.4128, 24.1262,\n",
       "          0.2837, 10.8538]], device='cuda:0')"
      ]
     },
     "execution_count": 117,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ll_dist.scale ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dpddm",
   "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}