experiment_QAM_constellation_classification.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-04-27 18:53:40.675886: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import truncnorm\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Random seed setting\n",
    "TRAIN_VAL_TEST_CLASS_SAMPLER_RANDOM_SEED = 123\n",
    "\n",
    "\n",
    "class MetaDataset:\n",
    "    def __init__(self, random_state=None):\n",
    "        if random_state is None:\n",
    "            self.random_state = np.random\n",
    "        else:\n",
    "            self.random_state = random_state\n",
    "\n",
    "    def generate_meta_train_data(self, n_tasks: int, n_samples: int) -> list:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def generate_meta_test_data(self, n_tasks: int, n_samples_context: int, n_samples_test: int) -> list:\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# QAM data generating code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "class QAM16ToyMetaDataset():\n",
    "    constellation = np.array([[-3, 3], [-1, 3], [1, 3], [3, 3],\n",
    "                              [-3, 1], [-1, 1], [1, 1], [3, 1],\n",
    "                              [-3, -1], [-1, -1], [1, -1], [3, -1],\n",
    "                              [-3, -3], [-1, -3], [1, -3], [3, -3]])\n",
    "\n",
    "    def __init__(self, random_state=None):\n",
    "        self.random_state = np.random.RandomState(seed=random_state)\n",
    "\n",
    "    def generate_meta_train_data(self, n_tasks=10, n_samples=100):\n",
    "        \"\"\"\n",
    "        Generates meta training data (16 QAM constellation points)\n",
    "        Args:\n",
    "            n_tasks (int): number of tasks\n",
    "            n_samples (int): number of samples per task\n",
    "        Returns:\n",
    "            List of 2-tuples [(x_train_1, y_train_1), ... (x_train_n, y_train_n)] wherein x_train is a stack\n",
    "            of 2D coordinate arrays with shape (n_samples, 2) and y_train a stack of one-hot encodings with shape\n",
    "            (n_samples, 16)\n",
    "        \"\"\"\n",
    "        assert n_tasks > 0 or n_samples > 0\n",
    "\n",
    "        task_tuples = []\n",
    "        for _ in range(n_tasks):\n",
    "            points = self.constellation[self.random_state.choice(16, size=n_samples)]\n",
    "            one_hot = np.zeros((n_samples, 16))\n",
    "            for i in range(n_samples):\n",
    "                one_hot[i, np.where(np.all(points[i] == self.constellation, axis=1))[0][0]] = 1\n",
    "            task_tuples.append((points, one_hot))\n",
    "\n",
    "        assert len(task_tuples) == n_tasks\n",
    "        assert all([x_train.shape[0] == y_train.shape[0] == n_samples for x_train, y_train in task_tuples])\n",
    "\n",
    "        return task_tuples\n",
    "\n",
    "    def generate_meta_test_data(self, n_tasks=10, n_samples_context=100, n_samples_test=500):\n",
    "        \"\"\"\n",
    "        Generates meta test data (16 QAM constellation points)\n",
    "        Args:\n",
    "            n_tasks (int): number of tasks\n",
    "            n_samples_context (int): number of context samples per task\n",
    "            n_samples_test (int): number of test samples per task\n",
    "        Returns:\n",
    "            List of n_tasks 4-tuples [(x_context_1, y_context_1, x_test_1, y_test_1), ... ]\n",
    "            wherein x_context and x_test are stacks of 2D coordinate arrays with shape (n_samples_context, 2) and\n",
    "            (n_samples_test, 2) respectively, and y_context and y_test are stacks of one-hot encodings with shape\n",
    "            (n_samples_context, 16) and (n_samples_test, 16) respectively.\n",
    "        \"\"\"\n",
    "        assert n_tasks > 0 and n_samples_context > 0 and n_samples_test > 0\n",
    "\n",
    "        task_tuples = []\n",
    "        for _ in range(n_tasks):\n",
    "            context_points = self.constellation[self.random_state.choice(16, size=n_samples_context)]\n",
    "            context_one_hot = np.zeros((n_samples_context, 16))\n",
    "            for i in range(n_samples_context):\n",
    "                context_one_hot[i, np.where(np.all(context_points[i] == self.constellation, axis=1))[0][0]] = 1\n",
    "\n",
    "            test_points = self.constellation[self.random_state.choice(16, size=n_samples_test)]\n",
    "            test_one_hot = np.zeros((n_samples_test, 16))\n",
    "            for i in range(n_samples_test):\n",
    "                test_one_hot[i, np.where(np.all(test_points[i] == self.constellation, axis=1))[0][0]] = 1\n",
    "\n",
    "            task_tuples.append((context_points, context_one_hot, test_points, test_one_hot))\n",
    "\n",
    "        assert len(task_tuples) == n_tasks\n",
    "        assert all([x_context.shape[0] == y_context.shape[0] == n_samples_context for x_context, y_context, x_test, y_test in task_tuples])\n",
    "        assert all([x_test.shape[0] == y_test.shape[0] == n_samples_test for x_context, y_context, x_test, y_test in task_tuples])\n",
    "\n",
    "        return task_tuples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = QAM16ToyMetaDataset()\n",
    "meta_train_data = env.generate_meta_train_data(n_tasks=5, n_samples=50)\n",
    "meta_test_data = env.generate_meta_test_data(n_tasks=5, n_samples_context=5, n_samples_test=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiIAAAGdCAYAAAAvwBgXAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAeuUlEQVR4nO3df5DU9X348dedeAcEbiv00FAOOEmbDrXafo1ETDFg8QfMxPhHnLTTUVDnigxmqmYcOSeR9NuaS6vTZEYdQ+kEGieOZpqhZFpsZFB+dAQ16MUffHEGy29iIBB34cA7YPf7h1/uG4Q79nA/9947Ho+Zzx8s79vPi50Pu092P/uhplQqlQIAIIHa1AMAAOcvIQIAJCNEAIBkhAgAkIwQAQCSESIAQDJCBABIRogAAMkMST1Ab4rFYuzduzdGjhwZNTU1qccBAMpQKpXi0KFDMXbs2Kit7f09j6oOkb1790ZTU1PqMQCAc7Br164YN25cr2uqOkRGjhwZER/9QRoaGhJPAwCUo1AoRFNTU/freG+qOkROfhzT0NAgRABggCnntAonqwIAyQgRACAZIQIAJCNEAIBkhAgAkIwQAQCSESIAQDJCBABIpqovaJaVg4e74i/++eXYd6grxoysi2f/+poYNaIu9VgMcF3Hi/H0hu2x4+CRmDBqeNw2dWLUDdH6fDL5I8fizmWvxt78hzE2NzR+MHdK5IZfmHosBoHDHx6P+557I3b+5miMv2hYfPerfxojhvZ/FtSUSqVSVnf+1FNPxVNPPRXbt2+PiIg/+qM/iocffjhmzZpV1s8XCoXI5XKRz+crdmXVq/5+Vew/3HXa7Y0j6uK1b1xfkX1w/mlbuTmWrN8Wxd/621RbE9EyrTlaZ09ONxgD2hcffTF2HDh62u0TRg+LtQ9cl2AiBoubn1gfb+4unHb75eMa4qf3TPvE99+X1+9M/7k2bty4+M53vhObNm2Kn//853HdddfFl7/85XjnnXey3G2PeoqQiIj9h7viqr9f1c8TMRi0rdwci9edGiEREcVSxOJ126Jt5eY0gzGg9RQhERE7DhyNLz76Yj9PxGDRU4RERLy5uxA3P7G+X+fJNES+9KUvxezZs+P3f//34w/+4A/ikUceiREjRsTGjRuz3O0ZHTzc1WOEnLT/cFccPMsa+G1dx4uxZP22XtcsWb8tuo4X+2kiBoP8kWM9RshJOw4cjfyRY/00EYPF4Q+P9xghJ725uxCHPzzeTxP148mqJ06ciGeffTY6Ojpi6tSpZ1zT2dkZhULhlK1S/uKfX67oOoiIeHrD9tPeCfm4YumjdVCuO5e9WtF1cNJ9z71R0XWVkHmIvPXWWzFixIior6+Pu+++O5YvXx6TJ5/5M/O2trbI5XLdW1NTU8Xm2HeovHc6yl0HERE7Dh6p6DqIiNib/7Ci6+Cknb/p/Z22vq6rhMxD5LOf/Wy0t7fHK6+8EvPnz485c+bE5s1n/sy8tbU18vl897Zr166KzTFmZHnfiil3HURETBg1vKLrICJibG5oRdfBSeMvGlbRdZWQ6bdmzmTmzJkxadKkWLx48VnXVvJbMwcPd8X/KuNk1Ne/cb2v8lK2ruPF+MNvPt/rxzO1NRFb/m6Wr/JStvyRY3HF/37hrOt+8fANvspLnxz+8Hhc9q2fnXXd29+68RN9lbdqvjVzJsViMTo7O/t7tzFqRF00niUwGkfUiRD6pG5IbbRMa+51Tcu0ZhFCn+SGXxgTRvf+L9IJo4eJEPpsxNAhcfm43sPg8nEN/Xo9kUyfHVtbW2PdunWxffv2eOutt6K1tTXWrFkTf/VXf5Xlbnv02jeu7zFGXEeEc9U6e3LMu7Y5amtOvb22JmLeta4jwrlZ+8B1PcaI64jwSfz0nmk9xkilriPSF5l+NHPXXXfF6tWr45e//GXkcrm4/PLL48EHH4zrry/vBT+LC5pFuLIq2XBlVbLgyqpkJcsrq/bl9bvfzxHpi6xCBADITlWfIwIAcJIQAQCSESIAQDJCBABIRogAAMkIEQAgGSECACQjRACAZIQIAJCMEAEAkhEiAEAyQgQASEaIAADJCBEAIBkhAgAkI0QAgGSECACQjBABAJIRIgBAMkIEAEhGiAAAyQgRACAZIQIAJCNEAIBkhAgAkIwQAQCSESIAQDJCBABIRogAAMkIEQAgGSECACQjRACAZIQIAJCMEAEAkhEiAEAyQgQASEaIAADJCBEAIBkhAgAkI0QAgGSECACQjBABAJIRIgBAMkIEAEhGiAAAyWQaIm1tbXHVVVfFyJEjY8yYMXHLLbfEu+++m+UuAYABJNMQWbt2bSxYsCA2btwYq1atimPHjsUNN9wQHR0dWe4WABggakqlUqm/drZ///4YM2ZMrF27Nq699tqzri8UCpHL5SKfz0dDQ0M/TAgAfFJ9ef0e0k8zRUREPp+PiIhRo0ad8fc7Ozujs7Oz+9eFQqFf5gIA0ui3k1WLxWLce++98YUvfCEuu+yyM65pa2uLXC7XvTU1NfXXeABAAv320cz8+fPj+eefj//+7/+OcePGnXHNmd4RaWpq8tEMAAwgVffRzD333BP/8R//EevWresxQiIi6uvro76+vj9GAgCqQKYhUiqV4mtf+1osX7481qxZE83NzVnuDgAYYDINkQULFsQzzzwTK1asiJEjR8b7778fERG5XC6GDRuW5a4BgAEg03NEampqznj70qVLY+7cuWf9eV/fBYCBp2rOEenHS5QAAAOQ/2sGAEhGiAAAyQgRACAZIQIAJCNEAIBkhAgAkIwQAQCSESIAQDJCBABIRogAAMkIEQAgGSECACQjRACAZIQIAJCMEAEAkhEiAEAyQgQASEaIAADJCBEAIBkhAgAkI0QAgGSECACQjBABAJIRIgBAMkIEAEhGiAAAyQgRACAZIQIAJCNEAIBkhAgAkIwQAQCSESIAQDJCBABIRogAAMkIEQAgGSECACQjRACAZIQIAJCMEAEAkhEiAEAyQgQASEaIAADJCBEAIBkhAgAkMyT1ACkc7ToR3165ObYfOBITRw+Ph2ZPjmF1F6QeiwHOcUUWuo4X4+kN22PHwSMxYdTwuG3qxKgb4t+QfHL5I8fizmWvxt78hzE2NzR+MHdK5IZf2O9z1JRKpVJWd75u3bp49NFHY9OmTfHLX/4yli9fHrfcckvZP18oFCKXy0U+n4+GhoaKzNTyw9di1eZ9p91+/eQxseT2qyqyD84/jiuy0LZycyxZvy2Kv/UsXVsT0TKtOVpnT043GAPeFx99MXYcOHra7RNGD4u1D1z3ie+/L6/fmWZ1R0dHXHHFFfHkk09muZuy9fRiERGxavO+aPnha/08EYOB44ostK3cHIvXnRohERHFUsTidduibeXmNIMx4PUUIREROw4cjS8++mK/zpPpRzOzZs2KWbNmZbmLsh3tOtHji8VJqzbvi6NdJ7ydTtkcV2Sh63gxlqzf1uuaJeu3xddv+EMf09An+SPHeoyQk3YcOBr5I8f67WOaqjqCOzs7o1AonLJVyrfL/NdDuesgwnFFNp7esP20d0I+rlj6aB30xZ3LXq3oukqoqhBpa2uLXC7XvTU1NVXsvrcfOFLRdRDhuCIbOw6Wd7yUuw5O2pv/sKLrKqGqQqS1tTXy+Xz3tmvXrord98TRwyu6DiIcV2Rjwqjyjpdy18FJY3NDK7quEqoqROrr66OhoeGUrVIeKvMM83LXQYTjimzcNnVi1Nb0vqa25qN10Bc/mDulousqoapCJEvD6i6I6yeP6XXN9ZPHOKGQPnFckYW6IbXRMq251zUt05qdqEqf5YZfGBNGD+t1zYTRw/r1eiKZHsWHDx+O9vb2aG9vj4iIbdu2RXt7e+zcuTPL3fZoye1X9fii4XoPnCvHFVlonT055l3bfNo7I7U1EfOudR0Rzt3aB67rMUYqdR2Rvsj0gmZr1qyJGTNmnHb7nDlzYtmyZWf9+SwuaBbhCphkw3FFFlxZlaxkeWXVvrx+Zxoin1RWIQIAZKdqrqwKANAbIQIAJCNEAIBkhAgAkIwQAQCSESIAQDJCBABIRogAAMkIEQAgGSECACQjRACAZIQIAJCMEAEAkhEiAEAyQgQASEaIAADJCBEAIBkhAgAkI0QAgGSECACQjBABAJIRIgBAMkIEAEhGiAAAyQgRACAZIQIAJCNEAIBkhAgAkIwQAQCSESIAQDJCBABIRogAAMkIEQAgGSECACQjRACAZIQIAJCMEAEAkhEiAEAyQgQASEaIAADJCBEAIBkhAgAkI0QAgGSECACQjBABAJLplxB58sknY+LEiTF06ND4/Oc/H6+++mp/7BYAqHKZh8hzzz0X999/fyxatChef/31uOKKK+LGG2+Mffv2Zb1rAKDKZR4i//RP/xQtLS1xxx13xOTJk+P73/9+DB8+PH7wgx9kvWsAoMplGiJdXV2xadOmmDlz5v/fYW1tzJw5MzZs2HDa+s7OzigUCqdsAMDglWmI/PrXv44TJ07ExRdffMrtF198cbz//vunrW9ra4tcLte9NTU1ZTkeAJBYVX1rprW1NfL5fPe2a9eu1CMBABkakuWd/+7v/m5ccMEF8atf/eqU23/1q1/FJZdcctr6+vr6qK+vz3IkAKCKZPqOSF1dXVx55ZWxevXq7tuKxWKsXr06pk6dmuWuAYABINN3RCIi7r///pgzZ0587nOfiylTpsT3vve96OjoiDvuuCPrXQMAVS7zEPnqV78a+/fvj4cffjjef//9+JM/+ZP4r//6r9NOYAUAzj81pVKplHqInhQKhcjlcpHP56OhoSH1OABAGfry+l1V35oBAM4vQgQASEaIAADJCBEAIBkhAgAkI0QAgGSECACQjBABAJIRIgBAMkIEAEhGiAAAyQgRACAZIQIAJCNEAIBkhAgAkIwQAQCSESIAQDJCBABIRogAAMkIEQAgGSECACQjRACAZIQIAJCMEAEAkhEiAEAyQgQASEaIAADJCBEAIBkhAgAkI0QAgGSECACQjBABAJIRIgBAMkIEAEhGiAAAyQgRACAZIQIAJCNEAIBkhAgAkIwQAQCSESIAQDJCBABIRogAAMkIEQAgmSGpB0ih63gxnt6wPXYcPBITRg2P26ZOjLohmoxPxnFFFo52nYhvr9wc2w8ciYmjh8dDsyfHsLoLUo/FIFAtx1ZNqVQqZXHHjzzySPznf/5ntLe3R11dXXzwwQd9vo9CoRC5XC7y+Xw0NDRUZK62lZtjyfptUfytP3VtTUTLtOZonT25Ivvg/OO4IgstP3wtVm3ed9rt108eE0tuvyrBRAwWWR9bfXn9zuyfa11dXXHrrbfG/Pnzs9pFn7Wt3ByL1536YhERUSxFLF63LdpWbk4zGAOa44os9PRCERGxavO+aPnha/08EYNFtR1bmYXI3/7t38Z9990Xf/zHf5zVLvqk63gxlqzf1uuaJeu3RdfxYj9NxGDguCILR7tO9PhCcdKqzfviaNeJfpqIwaIaj62q+gC7s7MzCoXCKVulPL1h+2n/Yv24YumjdVAuxxVZ+HaZ76KVuw5OqsZjq6pCpK2tLXK5XPfW1NRUsfvecfBIRddBhOOKbGw/UN7xUu46OKkaj60+hcjChQujpqam123Lli3nPExra2vk8/nubdeuXed8Xx83YdTwiq6DCMcV2Zg4urzjpdx1cFI1Hlt9+vru17/+9Zg7d26vay699NJzHqa+vj7q6+vP+ed7c9vUifHIyv/T69votTUfrYNyOa7IwkOzJ8fTG3eWtQ76ohqPrT6FSGNjYzQ2NmY1S6bqhtRGy7TmWLyu5xMLW6Y1u+4DfeK4IgvD6i6I6yeP6fWkwusnj3E9EfqsGo+tzJ4dd+7cGe3t7bFz5844ceJEtLe3R3t7exw+fDirXZ5V6+zJMe/a5qitOfX22pqIede63gPnxnFFFpbcflVcP3nMGX/PdUT4JKrt2MrsgmZz586Nf/3Xfz3t9pdeeimmT59e1n1kcUGzCFfAJBuOK7JQLVe/ZPDJ8tjqy+t3ZiFSCVmFCACQnaq4sioAwNkIEQAgGSECACQjRACAZIQIAJCMEAEAkhEiAEAyQgQASEaIAADJCBEAIBkhAgAkI0QAgGSECACQjBABAJIRIgBAMkIEAEhGiAAAyQgRACAZIQIAJCNEAIBkhAgAkIwQAQCSESIAQDJCBABIRogAAMkIEQAgGSECACQjRACAZIQIAJCMEAEAkhEiAEAyQgQASEaIAADJCBEAIBkhAgAkI0QAgGSECACQjBABAJIRIgBAMkIEAEhGiAAAyQgRACAZIQIAJCNEAIBkhAgAkExmIbJ9+/a46667orm5OYYNGxaTJk2KRYsWRVdXV1a7BAAGmCFZ3fGWLVuiWCzG4sWL4zOf+Uy8/fbb0dLSEh0dHfHYY49ltVsAYACpKZVKpf7a2aOPPhpPPfVU/M///E9Z6wuFQuRyucjn89HQ0JDxdABAJfTl9Tuzd0TOJJ/Px6hRo3r8/c7Ozujs7Oz+daFQ6I+xAIBE+u1k1a1bt8bjjz8e8+bN63FNW1tb5HK57q2pqam/xgMAEuhziCxcuDBqamp63bZs2XLKz+zZsyduuummuPXWW6OlpaXH+25tbY18Pt+97dq1q+9/IgBgwOjzOSL79++PAwcO9Lrm0ksvjbq6uoiI2Lt3b0yfPj2uvvrqWLZsWdTWlt8+zhEBgIEn03NEGhsbo7Gxsay1e/bsiRkzZsSVV14ZS5cu7VOEAACDX2Ynq+7ZsyemT58eEyZMiMceeyz279/f/XuXXHJJVrsFAAaQzEJk1apVsXXr1ti6dWuMGzfulN/rx28MAwBVLLPPSubOnRulUumMGwBAhP9rBgBISIgAAMkIEQAgGSECACQjRACAZIQIAJCMEAEAkhEiAEAyQgQASEaIAADJCBEAIBkhAgAkI0QAgGSECACQjBABAJIRIgBAMkIEAEhGiAAAyQgRACAZIQIAJCNEAIBkhAgAkIwQAQCSESIAQDJCBABIRogAAMkIEQAgGSECACQjRACAZIQIAJCMEAEAkhEiAEAyQgQASEaIAADJCBEAIBkhAgAkI0QAgGSECACQjBABAJIRIgBAMkIEAEhGiAAAyQgRACAZIQIAJDMk9QApdB0vxtMbtseOg0diwqjhcdvUiVE3RJPxyRz+8Hjc99wbsfM3R2P8RcPiu1/90xgx9Lz8K0YFeb4iK9VybNWUSqVSVnd+8803R3t7e+zbty8uuuiimDlzZvzDP/xDjB07tqyfLxQKkcvlIp/PR0NDQ0Vmalu5OZas3xbF3/pT19ZEtExrjtbZkyuyD84/Nz+xPt7cXTjt9svHNcRP75mWYCIGA89XZCXrY6svr9+Zps+MGTPixz/+cbz77rvxk5/8JN577734yle+kuUue9W2cnMsXnfqAx8RUSxFLF63LdpWbk4zGANaTxESEfHm7kLc/MT6fp6IwcDzFVmptmMr0xC577774uqrr44JEybENddcEwsXLoyNGzfGsWPHstztGXUdL8aS9dt6XbNk/bboOl7sp4kYDA5/eLzHCDnpzd2FOPzh8X6aiMHA8xVZqcZjq98+DDp48GD86Ec/imuuuSYuvPDCM67p7OyMQqFwylYpT2/Yflr9fVyx9NE6KNd9z71R0XUQ4fmK7FTjsZV5iDz44IPxqU99KkaPHh07d+6MFStW9Li2ra0tcrlc99bU1FSxOXYcPFLRdRARsfM3Ryu6DiI8X5Gdajy2+hwiCxcujJqaml63LVu2dK9/4IEH4o033ogXXnghLrjggrj99tujp/NjW1tbI5/Pd2+7du069z/Zx0wYNbyi6yAiYvxFwyq6DiI8X5Gdajy2+vytmf3798eBAwd6XXPppZdGXV3dabfv3r07mpqa4uWXX46pU6eedV+V/NZM1/Fi/OE3n+/1LanamogtfzfLV+Mo2+EPj8dl3/rZWde9/a0bfZWXsnm+Iiv9dWz15fW7z8+MjY2N0djYeE6DFYsfnfzS2dl5Tj//SdQNqY2Wac2xeF3PJ+m0TGv2l5o+GTF0SFw+rqHXE1YvH9cgQugTz1dkpRqPrcz29Morr8QTTzwR7e3tsWPHjnjxxRfjL//yL2PSpEllvRuShdbZk2Petc1RW3Pq7bU1EfOu9b18zs1P75kWl487c/G7jgjnyvMVWam2YyuzC5q99dZb8Td/8zfxi1/8Ijo6OuLTn/503HTTTfGNb3wjfu/3fq+s+8jigmYR1XM1OQYXV1YlC56vyEqWx1ZfXr8zvbLqJ5VViAAA2amaK6sCAPRGiAAAyQgRACAZIQIAJCNEAIBkhAgAkIwQAQCSESIAQDJCBABIpqqvP33yoq+FQs//oRgAUF1Ovm6Xc/H2qg6RQ4cORUREU1NT4kkAgL46dOhQ5HK5XtdU9f81UywWY+/evTFy5Mioqak5+w/0QaFQiKampti1a5f/x+YsPFbl81iVz2NVPo9V33i8ypfVY1UqleLQoUMxduzYqK3t/SyQqn5HpLa2NsaNG5fpPhoaGhyoZfJYlc9jVT6PVfk8Vn3j8SpfFo/V2d4JOcnJqgBAMkIEAEjmvA2R+vr6WLRoUdTX16cepep5rMrnsSqfx6p8Hqu+8XiVrxoeq6o+WRUAGNzO23dEAID0hAgAkIwQAQCSESIAQDJCJCJuvvnmGD9+fAwdOjQ+/elPx2233RZ79+5NPVbV2b59e9x1113R3Nwcw4YNi0mTJsWiRYuiq6sr9WhV6ZFHHolrrrkmhg8fHr/zO7+Tepyq8+STT8bEiRNj6NCh8fnPfz5effXV1CNVnXXr1sWXvvSlGDt2bNTU1MS///u/px6parW1tcVVV10VI0eOjDFjxsQtt9wS7777buqxqtJTTz0Vl19+efdFzKZOnRrPP/98snmESETMmDEjfvzjH8e7774bP/nJT+K9996Lr3zlK6nHqjpbtmyJYrEYixcvjnfeeSe++93vxve///146KGHUo9Wlbq6uuLWW2+N+fPnpx6l6jz33HNx//33x6JFi+L111+PK664Im688cbYt29f6tGqSkdHR1xxxRXx5JNPph6l6q1duzYWLFgQGzdujFWrVsWxY8fihhtuiI6OjtSjVZ1x48bFd77zndi0aVP8/Oc/j+uuuy6+/OUvxzvvvJNmoBKnWbFiRammpqbU1dWVepSq94//+I+l5ubm1GNUtaVLl5ZyuVzqMarKlClTSgsWLOj+9YkTJ0pjx44ttbW1JZyqukVEafny5anHGDD27dtXiojS2rVrU48yIFx00UWlf/mXf0myb++IfMzBgwfjRz/6UVxzzTVx4YUXph6n6uXz+Rg1alTqMRhAurq6YtOmTTFz5szu22pra2PmzJmxYcOGhJMxmOTz+YgIz09nceLEiXj22Wejo6Mjpk6dmmQGIfL/PPjgg/GpT30qRo8eHTt37owVK1akHqnqbd26NR5//PGYN29e6lEYQH7961/HiRMn4uKLLz7l9osvvjjef//9RFMxmBSLxbj33nvjC1/4Qlx22WWpx6lKb731VowYMSLq6+vj7rvvjuXLl8fkyZOTzDJoQ2ThwoVRU1PT67Zly5bu9Q888EC88cYb8cILL8QFF1wQt99+e5TOk4vO9vWxiojYs2dP3HTTTXHrrbdGS0tLosn737k8VkD/WrBgQbz99tvx7LPPph6lan32s5+N9vb2eOWVV2L+/PkxZ86c2Lx5c5JZBu0l3vfv3x8HDhzodc2ll14adXV1p92+e/fuaGpqipdffjnZW1X9qa+P1d69e2P69Olx9dVXx7Jly6K2dtD27GnO5bhatmxZ3HvvvfHBBx9kPN3A0NXVFcOHD49/+7d/i1tuuaX79jlz5sQHH3zg3cge1NTUxPLly095zDjdPffcEytWrIh169ZFc3Nz6nEGjJkzZ8akSZNi8eLF/b7vIf2+x37S2NgYjY2N5/SzxWIxIiI6OzsrOVLV6stjtWfPnpgxY0ZceeWVsXTp0vMqQiI+2XHFR+rq6uLKK6+M1atXd7+oFovFWL16ddxzzz1ph2PAKpVK8bWvfS2WL18ea9asESF9VCwWk73mDdoQKdcrr7wSr732WvzZn/1ZXHTRRfHee+/FN7/5zZg0adJ58W5IX+zZsyemT58eEyZMiMceeyz279/f/XuXXHJJwsmq086dO+PgwYOxc+fOOHHiRLS3t0dExGc+85kYMWJE2uESu//++2POnDnxuc99LqZMmRLf+973oqOjI+64447Uo1WVw4cPx9atW7t/vW3btmhvb49Ro0bF+PHjE05WfRYsWBDPPPNMrFixIkaOHNl9vlEul4thw4Ylnq66tLa2xqxZs2L8+PFx6NCheOaZZ2LNmjXxs5/9LM1ASb6rU0XefPPN0owZM0qjRo0q1dfXlyZOnFi6++67S7t37049WtVZunRpKSLOuHG6OXPmnPGxeumll1KPVhUef/zx0vjx40t1dXWlKVOmlDZu3Jh6pKrz0ksvnfEYmjNnTurRqk5Pz01Lly5NPVrVufPOO0sTJkwo1dXVlRobG0t//ud/XnrhhReSzTNozxEBAKrf+fUBPwBQVYQIAJCMEAEAkhEiAEAyQgQASEaIAADJCBEAIBkhAgAkI0QAgGSECACQjBABAJIRIgBAMv8XkqTXmwT3xM8AAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "qam16_dataset = QAM16ToyMetaDataset(random_state=42)\n",
    "x_train, y_train = meta_train_data[0]\n",
    "plt.scatter(x_train[:, 0], x_train[:, 1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BNN for classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                   | 0/500 [00:02<?, ?it/s]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "in user code:\n\n    File \"/Users/klemens/Desktop/III_essy_coding/Experiment2_PACOH_NN/pacoh_nn/bnn/bnn_classification_svgd.py\", line 120, in step  *\n        score = tape.gradient(post_log_prob, self.particles)  # (k, n)\n    File \"/Users/klemens/Desktop/III_essy_coding/Experiment2_PACOH_NN/pacoh_nn/modules/batched_model.py\", line 48, in concat_and_vectorize_grads  *\n        return tf.reshape(vectorized_gradients, self._parameters_shape)\n\n    ValueError: Cannot reshape a tensor with 137120 elements to shape [10,13728] (137280 elements) for '{{node Reshape_100}} = Reshape[T=DT_FLOAT, Tshape=DT_INT32](concat, Reshape_100/shape)' with input shapes: [137120,1], [2] and with input tensors computed as partial shapes: input[1] = [10,13728].\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[9], line 8\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;66;03m# setting up and fitting the BNN\u001b[39;00m\n\u001b[1;32m      7\u001b[0m bnn \u001b[38;5;241m=\u001b[39m BayesianNeuralNetwork_classification_SVGD(x_context, y_context, num_classes \u001b[38;5;241m=\u001b[39m\u001b[38;5;241m16\u001b[39m, hidden_layer_sizes\u001b[38;5;241m=\u001b[39m(\u001b[38;5;241m64\u001b[39m, \u001b[38;5;241m64\u001b[39m, \u001b[38;5;241m64\u001b[39m, \u001b[38;5;241m64\u001b[39m), prior_weight\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0.001\u001b[39m, bandwidth\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1000.0\u001b[39m)\n\u001b[0;32m----> 8\u001b[0m \u001b[43mbnn\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx_val\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mx_test\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my_val\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43my_test\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_iter_fit\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m500\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlog_period\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m500\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Desktop/III_essy_coding/Experiment2_PACOH_NN/pacoh_nn/bnn/classification_algo.py:26\u001b[0m, in \u001b[0;36mClassificationModel.fit\u001b[0;34m(self, x_val, y_val, log_period, num_iter_fit)\u001b[0m\n\u001b[1;32m     24\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m pbar:\n\u001b[1;32m     25\u001b[0m     x_batch, y_batch \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mnext\u001b[39m(train_batch_sampler)\n\u001b[0;32m---> 26\u001b[0m     loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx_batch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my_batch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     27\u001b[0m     loss_list\u001b[38;5;241m.\u001b[39mappend(loss)\n\u001b[1;32m     29\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m i \u001b[38;5;241m%\u001b[39m log_period \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n",
      "File \u001b[0;32m~/Desktop/III_essy_coding/Experiment2_PACOH_NN/qam_exp/lib/python3.8/site-packages/tensorflow/python/util/traceback_utils.py:153\u001b[0m, in \u001b[0;36mfilter_traceback.<locals>.error_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    151\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    152\u001b[0m   filtered_tb \u001b[38;5;241m=\u001b[39m _process_traceback_frames(e\u001b[38;5;241m.\u001b[39m__traceback__)\n\u001b[0;32m--> 153\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m e\u001b[38;5;241m.\u001b[39mwith_traceback(filtered_tb) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    154\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    155\u001b[0m   \u001b[38;5;28;01mdel\u001b[39;00m filtered_tb\n",
      "File \u001b[0;32m/var/folders/v9/1l9jmkf14cg25_5d_j2jbyy80000gn/T/__autograph_generated_file97auc878.py:36\u001b[0m, in \u001b[0;36mouter_factory.<locals>.inner_factory.<locals>.tf__step\u001b[0;34m(self, x_batch, y_batch)\u001b[0m\n\u001b[1;32m     34\u001b[0m     prior_pre_factor \u001b[38;5;241m=\u001b[39m ag__\u001b[38;5;241m.\u001b[39mif_exp(ag__\u001b[38;5;241m.\u001b[39mld(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39msqrt_mode, (\u001b[38;5;28;01mlambda\u001b[39;00m : (\u001b[38;5;241m1\u001b[39m \u001b[38;5;241m/\u001b[39m ag__\u001b[38;5;241m.\u001b[39mconverted_call(ag__\u001b[38;5;241m.\u001b[39mld(math)\u001b[38;5;241m.\u001b[39msqrt, (ag__\u001b[38;5;241m.\u001b[39mld(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39mnum_train_samples,), \u001b[38;5;28;01mNone\u001b[39;00m, fscope))), (\u001b[38;5;28;01mlambda\u001b[39;00m : (\u001b[38;5;241m1\u001b[39m \u001b[38;5;241m/\u001b[39m ag__\u001b[38;5;241m.\u001b[39mld(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39mnum_train_samples)), \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mself.sqrt_mode\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m     35\u001b[0m     post_log_prob \u001b[38;5;241m=\u001b[39m (ag__\u001b[38;5;241m.\u001b[39mld(avg_log_likelihood) \u001b[38;5;241m+\u001b[39m (ag__\u001b[38;5;241m.\u001b[39mld(prior_pre_factor) \u001b[38;5;241m*\u001b[39m ag__\u001b[38;5;241m.\u001b[39mld(prior_prob)))\n\u001b[0;32m---> 36\u001b[0m score \u001b[38;5;241m=\u001b[39m ag__\u001b[38;5;241m.\u001b[39mconverted_call(ag__\u001b[38;5;241m.\u001b[39mld(tape)\u001b[38;5;241m.\u001b[39mgradient, (ag__\u001b[38;5;241m.\u001b[39mld(post_log_prob), ag__\u001b[38;5;241m.\u001b[39mld(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39mparticles), \u001b[38;5;28;01mNone\u001b[39;00m, fscope)\n\u001b[1;32m     37\u001b[0m particles_copy \u001b[38;5;241m=\u001b[39m ag__\u001b[38;5;241m.\u001b[39mconverted_call(ag__\u001b[38;5;241m.\u001b[39mld(tf)\u001b[38;5;241m.\u001b[39midentity, (ag__\u001b[38;5;241m.\u001b[39mld(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39mparticles,), \u001b[38;5;28;01mNone\u001b[39;00m, fscope)\n\u001b[1;32m     38\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m ag__\u001b[38;5;241m.\u001b[39mld(tf)\u001b[38;5;241m.\u001b[39mGradientTape() \u001b[38;5;28;01mas\u001b[39;00m tape:\n",
      "File \u001b[0;32m~/Desktop/III_essy_coding/Experiment2_PACOH_NN/pacoh_nn/modules/batched_model.py:72\u001b[0m, in \u001b[0;36mTFModuleBatched.call_parametrized.<locals>.grad_fn\u001b[0;34m(dy, variables)\u001b[0m\n\u001b[1;32m     70\u001b[0m     tampered_y \u001b[38;5;241m=\u001b[39m y \u001b[38;5;241m*\u001b[39m dy\n\u001b[1;32m     71\u001b[0m grads_x_w \u001b[38;5;241m=\u001b[39m tape\u001b[38;5;241m.\u001b[39mgradient(tampered_y, [x] \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mlist\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrainable_variables))\n\u001b[0;32m---> 72\u001b[0m grads_to_input \u001b[38;5;241m=\u001b[39m [grads_x_w[\u001b[38;5;241m0\u001b[39m], \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconcat_and_vectorize_grads\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgrads_x_w\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m]\n\u001b[1;32m     73\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m grads_to_input, [\u001b[38;5;28;01mNone\u001b[39;00m] \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mlen\u001b[39m(variables)\n",
      "File \u001b[0;32m/var/folders/v9/1l9jmkf14cg25_5d_j2jbyy80000gn/T/__autograph_generated_file2buth_aw.py:36\u001b[0m, in \u001b[0;36mouter_factory.<locals>.inner_factory.<locals>.tf__concat_and_vectorize_grads\u001b[0;34m(self, gradients)\u001b[0m\n\u001b[1;32m     34\u001b[0m         do_return \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m     35\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m\n\u001b[0;32m---> 36\u001b[0m ag__\u001b[38;5;241m.\u001b[39mif_stmt((ag__\u001b[38;5;241m.\u001b[39mld(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m_parameters_shape \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m), if_body, else_body, get_state, set_state, (\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdo_return\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mretval_\u001b[39m\u001b[38;5;124m'\u001b[39m), \u001b[38;5;241m2\u001b[39m)\n\u001b[1;32m     37\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m fscope\u001b[38;5;241m.\u001b[39mret(retval_, do_return)\n",
      "File \u001b[0;32m/var/folders/v9/1l9jmkf14cg25_5d_j2jbyy80000gn/T/__autograph_generated_file2buth_aw.py:32\u001b[0m, in \u001b[0;36mouter_factory.<locals>.inner_factory.<locals>.tf__concat_and_vectorize_grads.<locals>.else_body\u001b[0;34m()\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m     31\u001b[0m     do_return \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[0;32m---> 32\u001b[0m     retval_ \u001b[38;5;241m=\u001b[39m ag__\u001b[38;5;241m.\u001b[39mconverted_call(ag__\u001b[38;5;241m.\u001b[39mld(tf)\u001b[38;5;241m.\u001b[39mreshape, (ag__\u001b[38;5;241m.\u001b[39mld(vectorized_gradients), ag__\u001b[38;5;241m.\u001b[39mld(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m_parameters_shape), \u001b[38;5;28;01mNone\u001b[39;00m, fscope)\n\u001b[1;32m     33\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m     34\u001b[0m     do_return \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
      "\u001b[0;31mValueError\u001b[0m: in user code:\n\n    File \"/Users/klemens/Desktop/III_essy_coding/Experiment2_PACOH_NN/pacoh_nn/bnn/bnn_classification_svgd.py\", line 120, in step  *\n        score = tape.gradient(post_log_prob, self.particles)  # (k, n)\n    File \"/Users/klemens/Desktop/III_essy_coding/Experiment2_PACOH_NN/pacoh_nn/modules/batched_model.py\", line 48, in concat_and_vectorize_grads  *\n        return tf.reshape(vectorized_gradients, self._parameters_shape)\n\n    ValueError: Cannot reshape a tensor with 137120 elements to shape [10,13728] (137280 elements) for '{{node Reshape_100}} = Reshape[T=DT_FLOAT, Tshape=DT_INT32](concat, Reshape_100/shape)' with input shapes: [137120,1], [2] and with input tensors computed as partial shapes: input[1] = [10,13728].\n"
     ]
    }
   ],
   "source": [
    "from pacoh_nn.bnn import BayesianNeuralNetwork_classification_SVGD\n",
    "\n",
    "\n",
    "x_context, y_context, x_test, y_test = meta_test_data[0]\n",
    "\n",
    "# setting up and fitting the BNN\n",
    "bnn = BayesianNeuralNetwork_classification_SVGD(x_context, y_context, num_classes =16, hidden_layer_sizes=(64, 64, 64, 64), prior_weight=0.001, bandwidth=1000.0)\n",
    "bnn.fit(x_val=x_test, y_val=y_test, num_iter_fit=500, log_period=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:5 out of the last 5 calls to <function TFModuleBatched.get_variables_vectorized at 0x7f8ae6132670> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n",
      "WARNING:tensorflow:6 out of the last 6 calls to <function TFModuleBatched.get_variables_vectorized at 0x7f8b021ea700> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n",
      "Start meta-training -------------------- \n",
      "\n",
      "Iter 0/2000 - Time 47.54 sec - Train-Loss: -5.16210\n",
      "\n",
      "Iter 500/2000 - Time 12.97 sec - Train-Loss: -4.10003\n",
      "\tStart meta-test posterior inference in 1 batches ------------------\n",
      "\tMeta-Test batch #1 consisting of 5 tasks----\n",
      "WARNING:tensorflow:5 out of the last 7 calls to <function TFModuleBatched.get_variables_stacked_per_model at 0x7f8aecc3b160> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n",
      "WARNING:tensorflow:6 out of the last 9 calls to <function TFModuleBatched.get_variables_stacked_per_model at 0x7f8aecf60790> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n",
      "(15, 5, 16)\n",
      "tf.Tensor(\n",
      "[[[-6.3561897e+09  2.0766680e+07 -6.4360448e+09 ... -6.3309537e+09\n",
      "    6.2573082e+09 -6.4260132e+09]\n",
      "  [-2.2477676e+09  7.3438070e+06 -2.2760064e+09 ... -2.2388426e+09\n",
      "    2.2127992e+09 -2.2724595e+09]\n",
      "  [-1.3774059e+07  4.5001984e+04 -1.3947107e+07 ... -1.3719370e+07\n",
      "    1.3559779e+07 -1.3925367e+07]\n",
      "  [-4.2531210e+08  1.3895640e+06 -4.3065555e+08 ... -4.2362342e+08\n",
      "    4.1869565e+08 -4.2998426e+08]\n",
      "  [-1.3774059e+07  4.5001984e+04 -1.3947107e+07 ... -1.3719370e+07\n",
      "    1.3559779e+07 -1.3925367e+07]]\n",
      "\n",
      " [[-6.3188685e+09  2.1166890e+07 -6.3828669e+09 ... -6.2582031e+09\n",
      "    6.1803172e+09 -6.3375933e+09]\n",
      "  [-2.2531679e+09  7.5476440e+06 -2.2759882e+09 ... -2.2315359e+09\n",
      "    2.2037637e+09 -2.2598449e+09]\n",
      "  [-1.3915490e+07  4.6613902e+04 -1.4056428e+07 ... -1.3781890e+07\n",
      "    1.3610369e+07 -1.3956727e+07]\n",
      "  [-4.2269469e+08  1.4159435e+06 -4.2697578e+08 ... -4.1863654e+08\n",
      "    4.1342640e+08 -4.2394726e+08]\n",
      "  [-1.3915490e+07  4.6613902e+04 -1.4056428e+07 ... -1.3781890e+07\n",
      "    1.3610369e+07 -1.3956727e+07]]\n",
      "\n",
      " [[-6.3204700e+09  2.0914504e+07 -6.3516298e+09 ... -6.2957757e+09\n",
      "    6.2407629e+09 -6.3832274e+09]\n",
      "  [-2.2448650e+09  7.4282810e+06 -2.2559322e+09 ... -2.2360942e+09\n",
      "    2.2165548e+09 -2.2671547e+09]\n",
      "  [-1.3852798e+07  4.5839145e+04 -1.3921092e+07 ... -1.3798676e+07\n",
      "    1.3678099e+07 -1.3990346e+07]\n",
      "  [-4.3731293e+08  1.4470736e+06 -4.3946886e+08 ... -4.3560432e+08\n",
      "    4.3179792e+08 -4.4165514e+08]\n",
      "  [-1.3852798e+07  4.5839145e+04 -1.3921092e+07 ... -1.3798676e+07\n",
      "    1.3678099e+07 -1.3990346e+07]]\n",
      "\n",
      " ...\n",
      "\n",
      " [[-6.3806618e+09 -6.9008936e+07 -6.1748767e+09 ... -6.3664174e+09\n",
      "    6.2898975e+09 -6.2636078e+09]\n",
      "  [-2.2403971e+09 -2.4230626e+07 -2.1681411e+09 ... -2.2353953e+09\n",
      "    2.2085274e+09 -2.1992965e+09]\n",
      "  [-1.3730313e+07 -1.4849767e+05 -1.3287492e+07 ... -1.3699661e+07\n",
      "    1.3535001e+07 -1.3478429e+07]\n",
      "  [-4.5053750e+08 -4.8727090e+06 -4.3600710e+08 ... -4.4953178e+08\n",
      "    4.4412861e+08 -4.4227235e+08]\n",
      "  [-1.3730313e+07 -1.4849767e+05 -1.3287492e+07 ... -1.3699661e+07\n",
      "    1.3535001e+07 -1.3478429e+07]]\n",
      "\n",
      " [[-6.3628836e+09 -6.9214592e+07 -6.2258463e+09 ... -6.3624791e+09\n",
      "    6.2833070e+09 -6.2665385e+09]\n",
      "  [-2.2317870e+09 -2.4277078e+07 -2.1837207e+09 ... -2.2316449e+09\n",
      "    2.2038751e+09 -2.1979937e+09]\n",
      "  [-1.3610508e+07 -1.4805334e+05 -1.3317378e+07 ... -1.3609646e+07\n",
      "    1.3440291e+07 -1.3404423e+07]\n",
      "  [-4.3137277e+08 -4.6924130e+06 -4.2208221e+08 ... -4.3134528e+08\n",
      "    4.2597773e+08 -4.2484096e+08]\n",
      "  [-1.3610508e+07 -1.4805334e+05 -1.3317378e+07 ... -1.3609646e+07\n",
      "    1.3440291e+07 -1.3404423e+07]]\n",
      "\n",
      " [[-6.3727350e+09 -6.9241240e+07 -6.2088371e+09 ... -6.3769564e+09\n",
      "    6.2768026e+09 -6.2973138e+09]\n",
      "  [-2.2319037e+09 -2.4250152e+07 -2.1745032e+09 ... -2.2333832e+09\n",
      "    2.1983063e+09 -2.2054897e+09]\n",
      "  [-1.3623836e+07 -1.4802605e+05 -1.3273453e+07 ... -1.3632865e+07\n",
      "    1.3418752e+07 -1.3462600e+07]\n",
      "  [-4.5895235e+08 -4.9866220e+06 -4.4714877e+08 ... -4.5925638e+08\n",
      "    4.5204355e+08 -4.5352067e+08]\n",
      "  [-1.3623836e+07 -1.4802605e+05 -1.3273453e+07 ... -1.3632865e+07\n",
      "    1.3418752e+07 -1.3462600e+07]]], shape=(15, 5, 16), dtype=float32)\n",
      "tfp.distributions.Categorical(\"Categorical\", batch_shape=[5], event_shape=[], dtype=int32)\n",
      "[<tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>, <tfp.distributions.Independent 'IndependentOneHotCategorical' batch_shape=[] event_shape=[5, 16] dtype=int32>]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "components[0] batch shape must be compatible with cat shape and other component batch shapes ((5,) vs ())",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[8], line 8\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mpacoh_nn\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpacoh_nn_classification\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m PACOH_NN_Classification\n\u001b[1;32m      2\u001b[0m pacoh_model \u001b[38;5;241m=\u001b[39m PACOH_NN_Classification(meta_train_data, random_seed\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m22\u001b[39m, num_classes\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m16\u001b[39m, \n\u001b[1;32m      3\u001b[0m                                       num_iter_meta_train\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2000\u001b[39m, num_iter_meta_test\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2000\u001b[39m,\n\u001b[1;32m      4\u001b[0m                                   learn_likelihood\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m, hyper_prior_weight\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1e-4\u001b[39m, lr\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2e-5\u001b[39m)\n\u001b[0;32m----> 8\u001b[0m \u001b[43mpacoh_model\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmeta_fit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmeta_test_data\u001b[49m\u001b[43m[\u001b[49m\u001b[43m:\u001b[49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meval_period\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1000\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlog_period\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m500\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Desktop/III_essy_coding/Experiment2_PACOH_NN/pacoh_nn/pacoh_nn_classification.py:127\u001b[0m, in \u001b[0;36mPACOH_NN_Classification.meta_fit\u001b[0;34m(self, meta_val_data, log_period, eval_period, plot_period)\u001b[0m\n\u001b[1;32m    125\u001b[0m \u001b[38;5;66;03m# run validation and print results\u001b[39;00m\n\u001b[1;32m    126\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m meta_val_data \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28miter\u001b[39m \u001b[38;5;241m%\u001b[39m eval_period \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28miter\u001b[39m \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[0;32m--> 127\u001b[0m     eval_metrics_mean, eval_metrics_std \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmeta_eval_datasets\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmeta_val_data\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    128\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m key \u001b[38;5;129;01min\u001b[39;00m eval_metrics_mean:\n\u001b[1;32m    129\u001b[0m         message \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m- Val-\u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[38;5;124m: \u001b[39m\u001b[38;5;132;01m%.3f\u001b[39;00m\u001b[38;5;124m +- \u001b[39m\u001b[38;5;132;01m%.3f\u001b[39;00m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m%\u001b[39m (key, eval_metrics_mean[key], eval_metrics_std[key])\n",
      "File \u001b[0;32m~/Desktop/III_essy_coding/Experiment2_PACOH_NN/pacoh_nn/pacoh_nn_classification.py:169\u001b[0m, in \u001b[0;36mPACOH_NN_Classification.meta_eval_datasets\u001b[0;34m(self, meta_valid_data, max_tasks_parallel)\u001b[0m\n\u001b[1;32m    166\u001b[0m     \u001b[38;5;66;03m# perform training and evaluation\u001b[39;00m\n\u001b[1;32m    167\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_meta_test_training_loop(task_batch, eval_models, eval_models_step, log_period\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10000\u001b[39m,\n\u001b[1;32m    168\u001b[0m                                   num_iter\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_iter_meta_test, eval_period\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10000\u001b[39m)\n\u001b[0;32m--> 169\u001b[0m     _, _, eval_metrics_grouped \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_meta_test_models_eval\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtask_batch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meval_models\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    170\u001b[0m     eval_metrics_dict_per_task\u001b[38;5;241m.\u001b[39mappend(eval_metrics_grouped)\n\u001b[1;32m    172\u001b[0m eval_metrics_mean, eval_metrics_std \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_aggregate_eval_metrics_across_tasks(eval_metrics_dict_per_task)\n",
      "File \u001b[0;32m~/Desktop/III_essy_coding/Experiment2_PACOH_NN/pacoh_nn/meta_algo.py:77\u001b[0m, in \u001b[0;36mMetaLearner._meta_test_models_eval\u001b[0;34m(self, test_tasks, eval_models)\u001b[0m\n\u001b[1;32m     75\u001b[0m eval_metrics_grouped \u001b[38;5;241m=\u001b[39m {}\n\u001b[1;32m     76\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m task, model \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(test_tasks, eval_models):\n\u001b[0;32m---> 77\u001b[0m     res \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43meval\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mtask\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mval_data\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     78\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m key \u001b[38;5;129;01min\u001b[39;00m res:\n\u001b[1;32m     79\u001b[0m         l \u001b[38;5;241m=\u001b[39m eval_metrics_grouped\u001b[38;5;241m.\u001b[39mget(key, [])\n",
      "File \u001b[0;32m~/Desktop/III_essy_coding/Experiment2_PACOH_NN/pacoh_nn/bnn/classification_algo.py:40\u001b[0m, in \u001b[0;36mClassificationModel.eval\u001b[0;34m(self, x, y)\u001b[0m\n\u001b[1;32m     38\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21meval\u001b[39m(\u001b[38;5;28mself\u001b[39m, x, y):\n\u001b[1;32m     39\u001b[0m     x, y \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_handle_input_data(x, y, convert_to_tensor\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m---> 40\u001b[0m     _, pred_dist \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpredict\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     41\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlikelihood\u001b[38;5;241m.\u001b[39mcalculate_eval_metrics(pred_dist, y)\n",
      "File \u001b[0;32m~/Desktop/III_essy_coding/Experiment2_PACOH_NN/pacoh_nn/bnn/bnn_classification_svgd.py:87\u001b[0m, in \u001b[0;36mBayesianNeuralNetwork_classification_SVGD.predict\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     84\u001b[0m y_pred \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnn\u001b[38;5;241m.\u001b[39mcall_parametrized(x, nn_params)\n\u001b[1;32m     86\u001b[0m \u001b[38;5;66;03m# form mixture of predictive distributions\u001b[39;00m\n\u001b[0;32m---> 87\u001b[0m pred_dist \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlikelihood\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_pred_mixture_dist\u001b[49m\u001b[43m(\u001b[49m\u001b[43my_pred\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     89\u001b[0m \u001b[38;5;66;03m# unnormalize preds\u001b[39;00m\n\u001b[1;32m     90\u001b[0m y_pred \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_unnormalize_preds(y_pred)\n",
      "File \u001b[0;32m~/Desktop/III_essy_coding/Experiment2_PACOH_NN/pacoh_nn/modules/classification_likelihood.py:143\u001b[0m, in \u001b[0;36mCategorical_softmax_Likelihood.get_pred_mixture_dist\u001b[0;34m(self, y_pred, std)\u001b[0m\n\u001b[1;32m    141\u001b[0m \u001b[38;5;28mprint\u001b[39m(components)\n\u001b[1;32m    142\u001b[0m \u001b[38;5;66;03m#print(categorical)\u001b[39;00m\n\u001b[0;32m--> 143\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtfp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdistributions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mMixture\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcategorical\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcomponents\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mpredictive_mixture\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Desktop/III_essy_coding/Experiment2_PACOH_NN/qam_exp/lib/python3.8/site-packages/decorator.py:232\u001b[0m, in \u001b[0;36mdecorate.<locals>.fun\u001b[0;34m(*args, **kw)\u001b[0m\n\u001b[1;32m    230\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m kwsyntax:\n\u001b[1;32m    231\u001b[0m     args, kw \u001b[38;5;241m=\u001b[39m fix(args, kw, sig)\n\u001b[0;32m--> 232\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcaller\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mextras\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkw\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Desktop/III_essy_coding/Experiment2_PACOH_NN/qam_exp/lib/python3.8/site-packages/tensorflow_probability/python/distributions/distribution.py:342\u001b[0m, in \u001b[0;36m_DistributionMeta.__new__.<locals>.wrapped_init\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    339\u001b[0m \u001b[38;5;66;03m# Note: if we ever want to have things set in `self` before `__init__` is\u001b[39;00m\n\u001b[1;32m    340\u001b[0m \u001b[38;5;66;03m# called, here is the place to do it.\u001b[39;00m\n\u001b[1;32m    341\u001b[0m self_\u001b[38;5;241m.\u001b[39m_parameters \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 342\u001b[0m \u001b[43mdefault_init\u001b[49m\u001b[43m(\u001b[49m\u001b[43mself_\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    343\u001b[0m \u001b[38;5;66;03m# Note: if we ever want to override things set in `self` by subclass\u001b[39;00m\n\u001b[1;32m    344\u001b[0m \u001b[38;5;66;03m# `__init__`, here is the place to do it.\u001b[39;00m\n\u001b[1;32m    345\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m self_\u001b[38;5;241m.\u001b[39m_parameters \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    346\u001b[0m   \u001b[38;5;66;03m# We prefer subclasses will set `parameters = dict(locals())` because\u001b[39;00m\n\u001b[1;32m    347\u001b[0m   \u001b[38;5;66;03m# this has nearly zero overhead. However, failing to do this, we will\u001b[39;00m\n\u001b[1;32m    348\u001b[0m   \u001b[38;5;66;03m# resolve the input arguments dynamically and only when needed.\u001b[39;00m\n",
      "File \u001b[0;32m~/Desktop/III_essy_coding/Experiment2_PACOH_NN/qam_exp/lib/python3.8/site-packages/tensorflow_probability/python/distributions/mixture.py:156\u001b[0m, in \u001b[0;36m_Mixture.__init__\u001b[0;34m(self, cat, components, validate_args, allow_nan_stats, name)\u001b[0m\n\u001b[1;32m    153\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m di, d \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(components):\n\u001b[1;32m    154\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m tensorshape_util\u001b[38;5;241m.\u001b[39mis_compatible_with(static_batch_shape,\n\u001b[1;32m    155\u001b[0m                                              d\u001b[38;5;241m.\u001b[39mbatch_shape):\n\u001b[0;32m--> 156\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    157\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcomponents[\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m] batch shape must be compatible with cat \u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m    158\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mshape and other component batch shapes (\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m vs \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m)\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[1;32m    159\u001b[0m             di, static_batch_shape, d\u001b[38;5;241m.\u001b[39mbatch_shape))\n\u001b[1;32m    160\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m tensorshape_util\u001b[38;5;241m.\u001b[39mis_compatible_with(static_event_shape,\n\u001b[1;32m    161\u001b[0m                                              d\u001b[38;5;241m.\u001b[39mevent_shape):\n\u001b[1;32m    162\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    163\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcomponents[\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m] event shape must be compatible with other \u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m    164\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcomponent event shapes (\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m vs \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m)\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[1;32m    165\u001b[0m             di, static_event_shape, d\u001b[38;5;241m.\u001b[39mevent_shape))\n",
      "\u001b[0;31mValueError\u001b[0m: components[0] batch shape must be compatible with cat shape and other component batch shapes ((5,) vs ())"
     ]
    }
   ],
   "source": [
    "from pacoh_nn.pacoh_nn_classification import PACOH_NN_Classification\n",
    "pacoh_model = PACOH_NN_Classification(meta_train_data, random_seed=22, num_classes= 16, \n",
    "                                      num_iter_meta_train=2000, num_iter_meta_test=2000,\n",
    "                                  learn_likelihood=False, hyper_prior_weight=1e-4, lr=2e-5)\n",
    "\n",
    "\n",
    "\n",
    "pacoh_model.meta_fit(meta_test_data[:10], eval_period=1000, log_period=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "qam_exp",
   "language": "python",
   "name": "qam_exp"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}