init

.gitattributes

@@ -0,0 +1,33 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.cpkt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,12 @@
+---
+title: TabPFN
+emoji: 🐨
+colorFrom: gray
+colorTo: blue
+sdk: gradio
+sdk_version: 3.1.1
+app_file: app.py
+pinned: true
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

TabPFN/PrepareDatasets.ipynb

@@ -0,0 +1,373 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "import openml\n",
+    "import pandas as pd"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from tqdm import tqdm\n",
+    "\n",
+    "from datasets import load_openml_list, test_dids_classification, valid_large_classification, open_cc_dids, open_cc_valid_dids\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "The autoreload extension is already loaded. To reload it, use:\n",
+      "  %reload_ext autoreload\n"
+     ]
+    }
+   ],
+   "source": [
+    "%load_ext autoreload\n",
+    "\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "### Prepare test datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "renamer = {'name': 'Name', 'NumberOfFeatures': '# Features', 'NumberOfSymbolicFeatures': '# Categorical Features', 'NumberOfInstances': '# Instances', 'NumberOfMissingValues': '# NaNs', 'NumberOfClasses': '# Classes', 'MinorityClassSize': 'Minority Class Size'}\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "OrderedDict([(99,\n",
+       "              {'id': 99,\n",
+       "               'alias': 'OpenML-CC18',\n",
+       "               'main_entity_type': 'task',\n",
+       "               'name': 'OpenML-CC18 Curated Classification benchmark',\n",
+       "               'status': 'active',\n",
+       "               'creation_date': '2019-02-21 18:47:13',\n",
+       "               'creator': 1}),\n",
+       "             (225,\n",
+       "              {'id': 225,\n",
+       "               'alias': 'OpenML-friendly',\n",
+       "               'main_entity_type': 'task',\n",
+       "               'name': 'OpenML100-friendly',\n",
+       "               'status': 'active',\n",
+       "               'creation_date': '2019-09-16 19:41:46',\n",
+       "               'creator': 1})])"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "openml.study.list_suites()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "suite = openml.study.get_suite(suite_id=99)\n",
+    "tasks = openml.tasks.list_tasks(output_format=\"dataframe\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Using ``@`` in `pd.DataFrame.query <\n",
+    "# https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.query.html>`_\n",
+    "# accesses variables outside of the current dataframe.\n",
+    "tasks = tasks.query(\"tid in @suite.tasks\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tids = list(tasks[np.logical_and(np.logical_and((tasks.NumberOfInstances <= 2000), (tasks.NumberOfFeatures <= 100))\n",
+    "                                 , (tasks.NumberOfClasses <= 10))].tid)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "30"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "len(tids)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tids = list(tasks[tasks.NumberOfInstances <= 2000].tid)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "open_cc_dids = [openml.tasks.get_task(task_id).get_dataset().id for task_id in tids]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "open_ml_datasets, open_ml_datasets_df = load_openml_list(test_dids_classification, multiclass=True, shuffled=True, filter_for_nan=False, max_samples = 100000, num_feats=100, return_capped=True)\n"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "open_ml_datasets_df = open_ml_datasets_df[open_ml_datasets_df.NumberOfInstances > 10000]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\\begin{tabular}{lrrrrrrr}\n",
+      "\\toprule\n",
+      "                                  Name &  \\# Features &  \\# Categorical Features &  \\# Instances &  \\# Classes &  \\# NaNs &  Minority Class Size &    id \\\\\n",
+      "\\midrule\n",
+      "                    KDDCup09\\_appetency &         231 &                      39 &        50000 &          2 & 8024152 &                  890 &  1111 \\\\\n",
+      "                              airlines &           8 &                       5 &       539383 &          2 &       0 &               240264 &  1169 \\\\\n",
+      "                        bank-marketing &          17 &                      10 &        45211 &          2 &       0 &                 5289 &  1461 \\\\\n",
+      "                                 nomao &         119 &                      30 &        34465 &          2 &       0 &                 9844 &  1486 \\\\\n",
+      "                                 adult &          15 &                       9 &        48842 &          2 &    6465 &                11687 &  1590 \\\\\n",
+      "                             covertype &          55 &                      45 &       581012 &          7 &       0 &                 2747 &  1596 \\\\\n",
+      "                           numerai28.6 &          22 &                       1 &        96320 &          2 &       0 &                47662 & 23517 \\\\\n",
+      "                             connect-4 &          43 &                      43 &        67557 &          3 &       0 &                 6449 & 40668 \\\\\n",
+      "jungle\\_chess\\_2pcs\\_raw\\_endgame\\_complete &           7 &                       1 &        44819 &          3 &       0 &                 4335 & 41027 \\\\\n",
+      "                            APSFailure &         171 &                       1 &        76000 &          2 & 1078695 &                 1375 & 41138 \\\\\n",
+      "                                albert &          79 &                      53 &       425240 &          2 & 2734000 &               212620 & 41147 \\\\\n",
+      "                             MiniBooNE &          51 &                       1 &       130064 &          2 &       0 &                36499 & 41150 \\\\\n",
+      "                             guillermo &        4297 &                       1 &        20000 &          2 &       0 &                 8003 & 41159 \\\\\n",
+      "                              riccardo &        4297 &                       1 &        20000 &          2 &       0 &                 5000 & 41161 \\\\\n",
+      "                               volkert &         181 &                       1 &        58310 &         10 &       0 &                 1361 & 41166 \\\\\n",
+      "                                dionis &          61 &                       1 &       416188 &        355 &       0 &                  878 & 41167 \\\\\n",
+      "                                jannis &          55 &                       1 &        83733 &          4 &       0 &                 1687 & 41168 \\\\\n",
+      "                                helena &          28 &                       1 &        65196 &        100 &       0 &                  111 & 41169 \\\\\n",
+      "\\bottomrule\n",
+      "\\end{tabular}\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "print_table = open_ml_datasets_df\n",
+    "print_table = print_table[['name', 'NumberOfFeatures', 'NumberOfSymbolicFeatures', 'NumberOfInstances', 'NumberOfClasses', 'NumberOfMissingValues', 'MinorityClassSize']].copy()\n",
+    "print_table['id'] = print_table.index\n",
+    "print_table[['NumberOfFeatures', 'NumberOfSymbolicFeatures', 'NumberOfInstances', 'NumberOfClasses', 'NumberOfMissingValues', 'MinorityClassSize']] = print_table[['NumberOfFeatures', 'NumberOfSymbolicFeatures', 'NumberOfInstances', 'NumberOfClasses', 'NumberOfMissingValues', 'MinorityClassSize']].astype(int)\n",
+    "print_table = print_table.rename(columns=renamer)\n",
+    "print(print_table.to_latex(index=False))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "### Prepare Validation datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "open_cc_datasets, open_cc_datasets_df = load_openml_list(open_cc_dids, multiclass=True, shuffled=True, filter_for_nan=False, max_samples = 2000, num_feats=100, return_capped=True)\n",
+    "\n",
+    "def extend_datasets(datasets, filtering = False):\n",
+    "    extended_datasets = {}\n",
+    "    i = 0\n",
+    "    for d in tqdm(datasets):\n",
+    "        if ((not 'NumberOfFeatures' in datasets[d])\n",
+    "                or (not 'NumberOfClasses' in datasets[d])\n",
+    "                or (not 'NumberOfInstances' in datasets[d])\n",
+    "                # or datasets[d]['NumberOfFeatures'] >= num_feats\n",
+    "                or datasets[d]['NumberOfClasses'] <= 0):\n",
+    "            print(datasets[d])\n",
+    "            continue\n",
+    "        ds = openml.datasets.get_dataset(d, download_data=False)\n",
+    "        if filtering and (datasets[d]['NumberOfInstances'] < 150\n",
+    "                          or datasets[d]['NumberOfInstances'] > 2000\n",
+    "                         or datasets[d]['NumberOfFeatures'] > 100\n",
+    "                         or datasets[d]['NumberOfClasses'] > 10):\n",
+    "            continue\n",
+    "        extended_datasets[d] = datasets[d]\n",
+    "        extended_datasets[d].update(ds.qualities)\n",
+    "    \n",
+    "    return extended_datasets\n",
+    "\n",
+    "# All datasets\n",
+    "openml_list = openml.datasets.list_datasets()\n",
+    "openml_list = pd.DataFrame.from_dict(openml_list, orient=\"index\")\n",
+    "\n",
+    "# Select only classification\n",
+    "openml_list = openml_list[~openml_list['MajorityClassSize'].isna()]\n",
+    "\n",
+    "# Remove duplicated datasets\n",
+    "duplicated = openml_list.duplicated(subset=['MajorityClassSize', 'MaxNominalAttDistinctValues', 'MinorityClassSize',\n",
+    "       'NumberOfClasses', 'NumberOfFeatures', 'NumberOfInstances',\n",
+    "       'NumberOfInstancesWithMissingValues', 'NumberOfMissingValues',\n",
+    "       'NumberOfNumericFeatures', 'NumberOfSymbolicFeatures'], keep='first')\n",
+    "openml_list = openml_list[~duplicated]\n",
+    "\n",
+    "duplicated = openml_list.duplicated(subset=['name'], keep='first')\n",
+    "openml_list = openml_list[~duplicated]\n",
+    "\n",
+    "# Filter out datasets that don't have meta information or Don't fulfill other criteria\n",
+    "openml_list = openml_list.to_dict(orient='index')\n",
+    "openml_list = pd.DataFrame.from_dict(extend_datasets(openml_list, filtering=True), orient=\"index\")\n",
+    "\n",
+    "# Filter out datasets in Open CC\n",
+    "openml_list = openml_list[~openml_list.name.apply(lambda x: x in test_datasets_multiclass_df.name.values)]\n",
+    "openml_list['CFI'] = openml_list.apply(lambda x: str(x.NumberOfClasses) + '_' + str(x.NumberOfFeatures) + '_' + str(x.NumberOfInstances), axis = 1)\n",
+    "test_datasets_multiclass_df['CFI'] = test_datasets_multiclass_df.apply(lambda x: str(x.NumberOfClasses) + '_' + str(x.NumberOfFeatures) + '_' + str(x.NumberOfInstances), axis = 1)\n",
+    "openml_list = openml_list[~openml_list.CFI.apply(lambda x: x in test_datasets_multiclass_df.CFI.values)]\n",
+    "\n",
+    "# Remove time series and artificial data\n",
+    "openml_list = openml_list[~openml_list.name.apply(lambda x: 'autoUniv' in x)]\n",
+    "openml_list = openml_list[~openml_list.name.apply(lambda x: 'fri_' in x)]\n",
+    "openml_list = openml_list[~openml_list.name.apply(lambda x: 'FOREX' in x)]\n",
+    "\n",
+    "# Remove datasets that overlapped with Open CC closely by name\n",
+    "openml_list = openml_list[~openml_list.name.apply(lambda x: 'ilpd' in x)]\n",
+    "openml_list = openml_list[~openml_list.name.apply(lambda x: 'car' in x)]\n",
+    "openml_list = openml_list[~openml_list.name.apply(lambda x: 'pc1' in x)]\n",
+    "\n",
+    "# Remove datasets that didn't load\n",
+    "openml_list = openml_list[~openml_list.did.apply(lambda x: x in {1065, 40589, 41496, 770, 43097, 43148, 43255, 43595, 43786, 41701})]\n",
+    "\n",
+    "# Remove class skew\n",
+    "openml_list = openml_list[(openml_list.MinorityClassSize / openml_list.MajorityClassSize) > 0.05]\n",
+    "openml_list = openml_list[openml_list.AutoCorrelation != 1]\n",
+    "\n",
+    "# Remove too easy\n",
+    "openml_list = openml_list[openml_list.CfsSubsetEval_DecisionStumpAUC != 1]"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print_table = openml_list\n",
+    "print_table = print_table[['name', 'NumberOfFeatures', 'NumberOfSymbolicFeatures', 'NumberOfInstances', 'NumberOfClasses', 'NumberOfMissingValues', 'MinorityClassSize']].copy()\n",
+    "print_table['id'] = print_table.index\n",
+    "print_table[['NumberOfFeatures', 'NumberOfSymbolicFeatures', 'NumberOfInstances', 'NumberOfClasses', 'NumberOfMissingValues', 'MinorityClassSize']] = print_table[['NumberOfFeatures', 'NumberOfSymbolicFeatures', 'NumberOfInstances', 'NumberOfClasses', 'NumberOfMissingValues', 'MinorityClassSize']].astype(int)\n",
+    "print_table = print_table.rename(columns=renamer)\n",
+    "print(print_table.to_latex(index=False))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

TabPFN/README.md

@@ -0,0 +1,23 @@
+# TabPFN
+
+## Installation
+```
+git clone git@github.com:automl/TabPFN.git
+cd TabPFN
+conda create -n TabPFN python=3.7
+conda activate TabPFN
+pip install -r requirements.txt
+```
+
+To run the autogluon baseline please create a separate environment and install autogluon==0.4.0, installation in the same environment as our other baselines is not possible.
+
+## Usage
+TrainingTuningAndPrediction: Train a TabPFN, Prior Tune and predict using a pretrained model.
+
+TabularEvaluationVisualization: Run Baselines and load Baseline and TabPFN Results for comparison and plotting.
+
+PrepareDatasets: Notebook used to inspect Datasets (Not needed to run baselines / TabPFN).
+
+SytheticGPAblation: Ablation experiments for Gaussian Process fitting with differentiable Hyper Parameters.
+
+

TabPFN/SyntheticGPAblation.ipynb

@@ -0,0 +1,392 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import time\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "import numpy as np\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from model_builder import get_model, get_default_spec, save_model, load_model\n",
+    "\n",
+    "from scripts.model_configs import *"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "# Setting params"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = 'cuda'\n",
+    "base_path = os.path.join('.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def train_function(config_sample, i, add_name=''):\n",
+    "    start_time = time.time()\n",
+    "    N_epochs_to_save = 50\n",
+    "    \n",
+    "    def save_callback(model, epoch):\n",
+    "        if not hasattr(model, 'last_saved_epoch'):\n",
+    "            model.last_saved_epoch = 0\n",
+    "        if ((time.time() - start_time) / (maximum_runtime * 60 / N_epochs_to_save)) > model.last_saved_epoch:\n",
+    "            print('Saving model..')\n",
+    "            config_sample['epoch_in_training'] = epoch\n",
+    "            save_model(model, base_path, f'models_diff/prior_diff_real_checkpoint{add_name}_n_{i}_epoch_{model.last_saved_epoch}.cpkt',\n",
+    "                           config_sample)\n",
+    "            model.last_saved_epoch = model.last_saved_epoch + 1 # TODO: Rename to checkpoint\n",
+    "    \n",
+    "    model = get_model(config_sample\n",
+    "                      , device\n",
+    "                      , should_train=True\n",
+    "                      , verbose=1\n",
+    "                      , epoch_callback = save_callback)\n",
+    "    \n",
+    "    return"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "heading_collapsed": true,
+    "tags": []
+   },
+   "source": [
+    "# Check synthetic data fitting"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "#### Workflow functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "hidden": true,
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "def generate_test_data(test_gp_params):\n",
+    "    # Generate test data\n",
+    "    config = {**test_gp_params}\n",
+    "\n",
+    "    config['verbose'] = False\n",
+    "    config['differentiable'] = False\n",
+    "    #config['bptt'] = config['bptt_in_training']\n",
+    "\n",
+    "    model_test_data = get_model(config, device, should_train=False, verbose=True)\n",
+    "    (hp_embedding, data, targets_), targets = next(iter(model_test_data[3]))\n",
+    "    (hp_embedding, data, targets_), targets = (hp_embedding, data.to(device), targets_.to(device)), targets.to(device)\n",
+    "    \n",
+    "    return (hp_embedding, data, targets_), targets\n",
+    "\n",
+    "def evaluate_hp_range(model, hparam_true, vary_hparam_ind, data, targets, eval_pos, plot_step_size):\n",
+    "    losses, hparams = [], []\n",
+    "    for l in np.arange(-1.74, 1.74, plot_step_size):\n",
+    "        hparam = [*hparam_true]\n",
+    "        hparam[vary_hparam_ind] = l\n",
+    "        hp_embedding_used = torch.tensor(hparam).to(device).float()\n",
+    "        with torch.inference_mode():\n",
+    "            outputs = torch.sigmoid(model[2]((hp_embedding_used.repeat(data.shape[1], 1), data, targets.float()), single_eval_pos=eval_pos)).squeeze(-1)\n",
+    "        \n",
+    "        loss = torch.nn.BCELoss()(outputs.flatten(), targets[eval_pos:].flatten()).detach().cpu()\n",
+    "        losses += [loss]\n",
+    "        hparam_real = [diff_hparams_f[i][1](hp) for i, hp in enumerate(hparam)]\n",
+    "        hparams += [hparam_real]\n",
+    "        \n",
+    "        print(loss, hparam_real, hparam, outputs.shape)\n",
+    "    return np.array(losses), np.array(hparams)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def differentiable_hparam_tuning_workflow(config_sample, hparam_label, batch_size=4, N_grad_steps=50, plot_step_size=0.1):\n",
+    "    test_gp_params = {\n",
+    "        \"lengthscale\": 1.0,\n",
+    "        #\"lengthscale_mean\": true_lengthscale,\n",
+    "        #\"lengthscale_std\": 0.5,\n",
+    "        \"noise\": 0.2,\n",
+    "        \"outputscale\": 1.0,\n",
+    "        'batch_size': batch_size\n",
+    "    }\n",
+    "    config_sample.update(test_gp_params)\n",
+    "    (hp_embedding, data, targets_), targets = generate_test_data(config_sample)\n",
+    "    hparam_true = [diff_hparams_f[i][0](test_gp_params[hp]) for i, hp in enumerate(diff_hparams_keys)]\n",
+    "    #hparam_true = [test_gp_params[hp] for i, hp in enumerate(diff_hparams_keys)]\n",
+    "\n",
+    "    for vary_hparam_ind, vary_hparam_name in hparam_label:\n",
+    "        print(vary_hparam_name)\n",
+    "\n",
+    "        losses, hparams = evaluate_hp_range(model, hparam_true, vary_hparam_ind, data, targets, eval_pos, plot_step_size=plot_step_size)\n",
+    "\n",
+    "        # TODO: Make only one parameter diffable\n",
+    "        hparam = torch.tensor([*hparam_true]).to(device).float()\n",
+    "        hparam[vary_hparam_ind] = hparam[vary_hparam_ind] + 0.1 #random.random() * 2 - 1\n",
+    "        hparam = torch.nn.Parameter(hparam, requires_grad=True)\n",
+    "        hparam_grad_mask = torch.zeros_like(hparam)\n",
+    "        hparam_grad_mask[vary_hparam_ind] = 1\n",
+    "\n",
+    "        optimizer = torch.optim.Adam([hparam], lr=0.1)\n",
+    "    \n",
+    "        for t in range(N_grad_steps):\n",
+    "            style = hparam.repeat(data.shape[1], 1)\n",
+    "            outputs = torch.sigmoid(model[2]((style, data, targets.float()), single_eval_pos=eval_pos)).squeeze(-1)\n",
+    "            loss = torch.nn.BCELoss()(outputs.flatten(), targets[eval_pos:].flatten())\n",
+    "            optimizer.zero_grad()\n",
+    "            loss.backward()\n",
+    "            with torch.no_grad():\n",
+    "                hparam.grad *= hparam_grad_mask\n",
+    "            optimizer.step()\n",
+    "            print('loss:', loss, 'hparams', diff_hparams_f[vary_hparam_ind][1](hparam[vary_hparam_ind]), 'true', diff_hparams_f[vary_hparam_ind][1](hparam_true[vary_hparam_ind]))\n",
+    "        inferred_param = diff_hparams_f[vary_hparam_ind][1](hparam[vary_hparam_ind].cpu().detach().numpy())\n",
+    "        return hparams, losses, inferred_param, vary_hparam_ind, hparam_true\n",
+    "        "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "#### Fitting a PFN with HP-Diffable GP Prior"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {
+    "hidden": true,
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "num_features = 5\n",
+    "bptt = 200\n",
+    "eval_positions = [100]\n",
+    "\n",
+    "config_general = get_general_config(num_features, bptt, eval_positions)\n",
+    "config_flexible_categorical = get_flexible_categorical_config(num_features)\n",
+    "\n",
+    "config_gp = {'noise': 0.2,  \"lengthscale\": 1.0, \"outputscale\": 1.0}\n",
+    "config_diff_gp = {'differentiable_hyperparameters': {\n",
+    "        'outputscale': {'distribution': 'uniform', 'min': 0., 'max': 10.0},\n",
+    "        'lengthscale': {'distribution': 'uniform', 'min': 0., 'max': 10.0},\n",
+    "        'noise': {'distribution': 'uniform', 'min': 0.0000001, 'max': 0.5},\n",
+    "    }\n",
+    "}\n",
+    "\n",
+    "config = {**config_general, **config_flexible_categorical, **config_diff_gp, **config_gp}\n",
+    "\n",
+    "config['prior_type'], config['differentiable'], config['flexible'] = 'gp', True, True\n",
+    "config['num_features'], config['num_features_used'] = num_features, num_features\n",
+    "config['epochs'], config['num_steps'], config['verbose'] = 500, 100, False\n",
+    "config[\"lr\"] = 0.00001\n",
+    "config[\"dropout\"] = 0\n",
+    "config[\"emsize\"] = 512\n",
+    "config[\"batch_size\"] = 128\n",
+    "config[\"aggregate_k_gradients\"] = 1\n",
+    "config['set_value_to_nan'] = 0.0\n",
+    "config['output_multiclass_ordered_p'] = 1.0\n",
+    "config['categorical_feature_p'] = 0.0\n",
+    "config['nan_prob_a_reason'] = 0.0\n",
+    "config['nan_prob_no_reason'] = 0.0\n",
+    "config['nan_prob_unknown_reason'] = 0.0\n",
+    "config[\"nlayers\"] = 8\n",
+    "\n",
+    "# TODO: This should not be sampled, but be one config\n",
+    "# TODO: This uses old hyperparam sampler throws error\n",
+    "config_sample = evaluate_hypers(config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {
+    "hidden": true,
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using style prior: True\n",
+      "Using cpu:0 device\n",
+      "Not using distributed\n",
+      "DataLoader.__dict__ {'num_steps': 100, 'fuse_x_y': False, 'get_batch_kwargs': {'batch_size': 128, 'seq_len': 200, 'seq_len_maximum': 200, 'device': 'cpu:0', 'num_features': 5, 'hyperparameters': {'lr': 1e-05, 'dropout': 0, 'emsize': 512, 'batch_size': 128, 'nlayers': 8, 'num_features': 5, 'nhead': 4, 'nhid_factor': 2, 'bptt': 200, 'eval_positions': None, 'seq_len_used': 200, 'sampling': 'normal', 'epochs': 500, 'num_steps': 100, 'verbose': False, 'pre_sample_causes': True, 'mix_activations': False, 'nan_prob_unknown_reason_reason_prior': 1.0, 'categorical_feature_p': 0.0, 'nan_prob_no_reason': 0.0, 'nan_prob_unknown_reason': 0.0, 'nan_prob_a_reason': 0.0, 'max_num_classes': 2, 'num_classes': 2, 'noise_type': 'Gaussian', 'balanced': True, 'normalize_to_ranking': False, 'set_value_to_nan': 0.0, 'normalize_by_used_features': True, 'num_features_used': 5, 'differentiable_hyperparameters': {'distribution': 'uniform', 'min': 0.0, 'max': 10.0}, 'noise': 0.2, 'lengthscale': 1.0, 'outputscale': 1.0, 'prior_type': 'gp', 'differentiable': True, 'flexible': True, 'aggregate_k_gradients': 1, 'output_multiclass_ordered_p': 1.0, 'recompute_attn': False}, 'num_outputs': 1, 'dynamic_batch_size': 2, 'get_batch': <function get_model.<locals>.make_get_batch.<locals>.<lambda> at 0x7f39ad8dcf80>, 'differentiable_hyperparameters': {'outputscale': {'distribution': 'uniform', 'min': 0.0, 'max': 10.0}, 'lengthscale': {'distribution': 'uniform', 'min': 0.0, 'max': 10.0}, 'noise': {'distribution': 'uniform', 'min': 1e-07, 'max': 0.5}}}, 'num_features': 5, 'num_outputs': 1}\n",
+      "Using a Transformer with 17.35 M parameters\n"
+     ]
+    }
+   ],
+   "source": [
+    "device = 'cuda'\n",
+    "train_function(config_sample, 0, add_name='gp_experiments_diff_with_noise_no_meta_new')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "#### Evaluating a PFN (with pretrained model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {
+    "hidden": true,
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using style prior: True\n",
+      "Using cpu:0 device\n",
+      "Not using distributed\n",
+      "DataLoader.__dict__ {'num_steps': 100, 'fuse_x_y': False, 'get_batch_kwargs': {'batch_size': 1, 'seq_len': 10, 'seq_len_maximum': 10, 'device': 'cpu:0', 'num_features': 5, 'hyperparameters': {'lr': 1e-05, 'dropout': 0, 'emsize': 512, 'batch_size': 1, 'nlayers': 8, 'num_features': 5, 'nhead': 4, 'nhid_factor': 2, 'bptt': 10, 'eval_positions': [190], 'seq_len_used': 200, 'sampling': 'normal', 'epochs': 500, 'num_steps': 100, 'verbose': False, 'pre_sample_causes': True, 'mix_activations': False, 'nan_prob_unknown_reason_reason_prior': 1.0, 'output_multiclass_ordered_p': 1.0, 'categorical_feature_p': 0.0, 'nan_prob_no_reason': 0.0, 'nan_prob_unknown_reason': 0.0, 'nan_prob_a_reason': 0.0, 'max_num_classes': 2, 'num_classes': 2, 'noise_type': 'Gaussian', 'balanced': True, 'multiclass_type': 'rank', 'normalize_to_ranking': False, 'set_value_to_nan': 0.0, 'normalize_by_used_features': True, 'num_features_used': <function load_model.<locals>.<lambda> at 0x7f39ad8534d0>, 'differentiable_hyperparameters': {'distribution': 'uniform', 'min': 0.0, 'max': 10.0}, 'noise': 0.03, 'lengthscale': 1.0, 'outputscale': 1.0, 'prior_type': 'gp', 'differentiable': True, 'flexible': True, 'aggregate_k_gradients': 1, 'recompute_attn': False, 'bptt_extra_samples': None, 'epoch_in_training': 0.998, 'categorical_features_sampler': <function load_model.<locals>.<lambda> at 0x7f39ad853680>, 'num_features_used_in_training': 5, 'num_classes_in_training': 2, 'batch_size_in_training': 128, 'bptt_in_training': 200, 'bptt_extra_samples_in_training': None}, 'num_outputs': 1, 'dynamic_batch_size': 2, 'get_batch': <function get_model.<locals>.make_get_batch.<locals>.<lambda> at 0x7f39ad81ab90>, 'differentiable_hyperparameters': {'outputscale': {'distribution': 'uniform', 'min': 0.0, 'max': 10.0}, 'lengthscale': {'distribution': 'uniform', 'min': 0.0, 'max': 10.0}, 'noise': {'distribution': 'uniform', 'min': 1e-07, 'max': 0.5}}}, 'num_features': 5, 'num_outputs': 1}\n",
+      "Using a Transformer with 17.35 M parameters\n"
+     ]
+    }
+   ],
+   "source": [
+    "device = 'cpu'\n",
+    "model, c = load_model(base_path, f'models_diff/gp_ablation_model.cpkt', device, eval_positions, verbose=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from priors.differentiable_prior import DifferentiableHyperparameterList\n",
+    "diff_list = DifferentiableHyperparameterList(c['differentiable_hyperparameters'], 512, device)\n",
+    "diff_hparams_keys, diff_hparams_f = diff_list.get_hyperparameter_info()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "model[2].eval()\n",
+    "eval_pos = 100\n",
+    "\n",
+    "hparam_label = [(1, 'outputscale')]\n",
+    "hparam_label = [(0, 'lengthscale')]\n",
+    "hparam_label = [(2, 'noise')]\n",
+    "hparam_labels = [[(1, 'outputscale')], [(2, 'noise')], [(0, 'lengthscale')]]\n",
+    "#hparam_labels = [[(2, 'noise')]]\n",
+    "\n",
+    "hparams, losses, inferred_param, vary_hparam_ind, hparam_true = {}, {}, {}, {}, {}\n",
+    "\n",
+    "for hparam_label in hparam_labels:\n",
+    "    (hparams[hparam_label[0][1]], losses[hparam_label[0][1]], inferred_param[hparam_label[0][1]], vary_hparam_ind[hparam_label[0][1]], \n",
+    "     hparam_true[hparam_label[0][1]]) = differentiable_hparam_tuning_workflow(config_sample, \n",
+    "                                                                                                          hparam_label=hparam_label, \n",
+    "                                                                                                          batch_size=256, \n",
+    "                                                                                                          N_grad_steps=50,\n",
+    "                                                                             plot_step_size=0.05)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "label = 'lengthscale'\n",
+    "\n",
+    "#import tikzplotlib\n",
+    "\n",
+    "inferred = losses[label]\n",
+    "\n",
+    "plt.plot(hparams[label][:, vary_hparam_ind[label]], losses[label])\n",
+    "true = diff_hparams_f[vary_hparam_ind[label]][1](hparam_true[label][vary_hparam_ind[label]])\n",
+    "plt.axvline(x=inferred_param[label], linestyle='solid', color='red')\n",
+    "plt.axvline(x=true, linestyle='dashed')\n",
+    "\n",
+    "plt.ylabel('Cross entropy Loss')\n",
+    "plt.xlabel(label)\n",
+    "\n",
+    "#tikzplotlib.save(f'diff_inferred_params_{label}.tex', axis_height='5.2cm', axis_width='5.2cm', strict=True)\n",
+    "\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

TabPFN/TabPFNPredictionOnly.ipynb

@@ -0,0 +1,253 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This notebook shows how to use TabPFN for tabular prediction with a scikit learn wrapper.\n",
+    "\n",
+    "classifier = TabPFNClassifier(device='cpu')\n",
+    "classifier.fit(train_xs, train_ys)\n",
+    "prediction_ = classifier.predict(test_xs)\n",
+    "\n",
+    "The fit function does not perform any computations, but only saves the training data. Computations are only done at inference time, when calling predict.\n",
+    "Note that the presaved models were trained for up to 100 features, 10 classes and 1000 samples. While the model does not have a hard bound on the number of samples, the features and classes are restricted and larger sizes lead to an error."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "### Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "import os\n",
+    "import random\n",
+    "\n",
+    "from model_builder import get_model, get_default_spec, save_model, load_model\n",
+    "from scripts.transformer_prediction_interface import transformer_predict, get_params_from_config, TabPFNClassifier\n",
+    "\n",
+    "from datasets import load_openml_list, open_cc_dids, open_cc_valid_dids\n",
+    "\n",
+    "from scripts import tabular_metrics"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "base_path = '.'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "### Load datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "outputs_hidden": true
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "max_samples = 10000\n",
+    "bptt = 10000\n",
+    "\n",
+    "cc_test_datasets_multiclass, cc_test_datasets_multiclass_df = load_openml_list(open_cc_dids, multiclass=True, shuffled=True, filter_for_nan=False, max_samples = max_samples, num_feats=100, return_capped=True)\n",
+    "cc_valid_datasets_multiclass, cc_valid_datasets_multiclass_df = load_openml_list(open_cc_valid_dids, multiclass=True, shuffled=True, filter_for_nan=False, max_samples = max_samples, num_feats=100, return_capped=True)\n",
+    "\n",
+    "# Loading longer OpenML Datasets for generalization experiments (optional)\n",
+    "# test_datasets_multiclass, test_datasets_multiclass_df = load_openml_list(test_dids_classification, multiclass=True, shuffled=True, filter_for_nan=False, max_samples = 10000, num_feats=100, return_capped=True)\n",
+    "\n",
+    "random.seed(0)\n",
+    "random.shuffle(cc_valid_datasets_multiclass)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import get_openml_classification"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset = openml.datasets.get_dataset(31)\n",
+    "X, y, categorical_indicator, attribute_names = dataset.get_data(\n",
+    "        dataset_format=\"array\", target=dataset.default_target_attribute\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_datasets(selector, task_type, suite='cc'):\n",
+    "    if task_type == 'binary':\n",
+    "        ds = valid_datasets_binary if selector == 'valid' else test_datasets_binary\n",
+    "    else:\n",
+    "        if suite == 'openml':\n",
+    "            ds = valid_datasets_multiclass if selector == 'valid' else test_datasets_multiclass\n",
+    "        elif suite == 'cc':\n",
+    "            ds = cc_valid_datasets_multiclass if selector == 'valid' else cc_test_datasets_multiclass\n",
+    "        else:\n",
+    "            raise Exception(\"Unknown suite\")\n",
+    "    return ds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_string, longer, task_type = '', 1, 'multiclass'\n",
+    "eval_positions = [1000]\n",
+    "bptt = 2000\n",
+    "    \n",
+    "test_datasets, valid_datasets = get_datasets('test', task_type, suite='cc'), get_datasets('valid', task_type, suite='cc')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "jp-MarkdownHeadingCollapsed": true,
+    "tags": []
+   },
+   "source": [
+    "### Select a dataset for prediction"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "[(i, test_datasets[i][0]) for i in range(len(test_datasets))]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "evaluation_dataset_index = 4 # Index of the dataset to predict\n",
+    "ds = test_datasets[evaluation_dataset_index]\n",
+    "print(f'Evaluation dataset name: {ds[0]} shape {ds[1].shape}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "xs, ys = ds[1].clone(), ds[2].clone()\n",
+    "eval_position = xs.shape[0] // 2\n",
+    "train_xs, train_ys = xs[0:eval_position], ys[0:eval_position]\n",
+    "test_xs, test_ys = xs[eval_position:], ys[eval_position:]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "### Predict using a Fitted and Tuned Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "classifier = TabPFNClassifier(device='cpu')\n",
+    "classifier.fit(train_xs, train_ys)\n",
+    "prediction_ = classifier.predict_proba(test_xs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "roc, ce = tabular_metrics.auc_metric(test_ys, prediction_), tabular_metrics.cross_entropy(test_ys, prediction_)\n",
+    "'AUC', float(roc), 'Cross Entropy', float(ce)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

TabPFN/datasets/__init__.py

@@ -0,0 +1,149 @@
+import pandas as pd
+import torch
+import numpy as np
+import openml
+
+
+def get_openml_classification(did, max_samples, multiclass=True, shuffled=True):
+    dataset = openml.datasets.get_dataset(did)
+    X, y, categorical_indicator, attribute_names = dataset.get_data(
+        dataset_format="array", target=dataset.default_target_attribute
+    )
+
+    if not multiclass:
+        X = X[y < 2]
+        y = y[y < 2]
+
+    if multiclass and not shuffled:
+        raise NotImplementedError("This combination of multiclass and shuffling isn't implemented")
+
+    if not isinstance(X, np.ndarray) or not isinstance(y, np.ndarray):
+        print('Not a NP Array, skipping')
+        return None, None, None, None
+
+    if not shuffled:
+        sort = np.argsort(y) if y.mean() < 0.5 else np.argsort(-y)
+        pos = int(y.sum()) if y.mean() < 0.5 else int((1 - y).sum())
+        X, y = X[sort][-pos * 2:], y[sort][-pos * 2:]
+        y = torch.tensor(y).reshape(2, -1).transpose(0, 1).reshape(-1).flip([0]).float()
+        X = torch.tensor(X).reshape(2, -1, X.shape[1]).transpose(0, 1).reshape(-1, X.shape[1]).flip([0]).float()
+    else:
+        order = np.arange(y.shape[0])
+        np.random.seed(13)
+        np.random.shuffle(order)
+        X, y = torch.tensor(X[order]), torch.tensor(y[order])
+    if max_samples:
+        X, y = X[:max_samples], y[:max_samples]
+
+    return X, y, list(np.where(categorical_indicator)[0]), attribute_names
+
+def load_openml_list(dids, filter_for_nan=False
+                     , num_feats=100
+                     , min_samples = 100
+                     , max_samples=400
+                     , multiclass=True
+                     , max_num_classes=10
+                     , shuffled=True
+                     , return_capped = False):
+    datasets = []
+    openml_list = openml.datasets.list_datasets(dids)
+    print(f'Number of datasets: {len(openml_list)}')
+
+    datalist = pd.DataFrame.from_dict(openml_list, orient="index")
+    if filter_for_nan:
+        datalist = datalist[datalist['NumberOfInstancesWithMissingValues'] == 0]
+        print(f'Number of datasets after Nan and feature number filtering: {len(datalist)}')
+
+    for ds in datalist.index:
+        modifications = {'samples_capped': False, 'classes_capped': False, 'feats_capped': False}
+        entry = datalist.loc[ds]
+
+        print('Loading', entry['name'], entry.did, '..')
+
+        if entry['NumberOfClasses'] == 0.0:
+            raise Exception("Regression not supported")
+            #X, y, categorical_feats, attribute_names = get_openml_regression(int(entry.did), max_samples)
+        else:
+            X, y, categorical_feats, attribute_names = get_openml_classification(int(entry.did), max_samples
+                                                                , multiclass=multiclass, shuffled=shuffled)
+        if X is None:
+            continue
+
+        if X.shape[1] > num_feats:
+            if return_capped:
+                X = X[:, 0:num_feats]
+                categorical_feats = [c for c in categorical_feats if c < num_feats]
+                modifications['feats_capped'] = True
+            else:
+                print('Too many features')
+                continue
+        if X.shape[0] == max_samples:
+            modifications['samples_capped'] = True
+
+        if X.shape[0] < min_samples:
+            print(f'Too few samples left')
+            continue
+
+        if len(np.unique(y)) > max_num_classes:
+            if return_capped:
+                X = X[y < np.unique(y)[10]]
+                y = y[y < np.unique(y)[10]]
+                modifications['classes_capped'] = True
+            else:
+                print(f'Too many classes')
+                continue
+
+        datasets += [[entry['name'], X, y, categorical_feats, attribute_names, modifications]]
+
+    return datasets, datalist
+
+
+# Classification
+valid_dids_classification = [13, 59, 4, 15, 40710, 43, 1498]
+test_dids_classification = [973, 1596, 40981, 1468, 40984, 40975, 41163, 41147, 1111, 41164, 1169, 1486, 41143, 1461, 41167, 40668, 41146, 41169, 41027, 23517, 41165, 41161, 41159, 41138, 1590, 41166, 1464, 41168, 41150, 1489, 41142, 3, 12, 31, 54, 1067]
+valid_large_classification = [  943, 23512,    49,   838,  1131,   767,  1142,   748,  1112,
+        1541,   384,   912,  1503,   796,    20,    30,   903,  4541,
+         961,   805,  1000,  4135,  1442,   816,  1130,   906,  1511,
+         184,   181,   137,  1452,  1481,   949,   449,    50,   913,
+        1071,   831,   843,     9,   896,  1532,   311,    39,   451,
+         463,   382,   778,   474,   737,  1162,  1538,   820,   188,
+         452,  1156,    37,   957,   911,  1508,  1054,   745,  1220,
+         763,   900,    25,   387,    38,   757,  1507,   396,  4153,
+         806,   779,   746,  1037,   871,   717,  1480,  1010,  1016,
+         981,  1547,  1002,  1126,  1459,   846,   837,  1042,   273,
+        1524,   375,  1018,  1531,  1458,  6332,  1546,  1129,   679,
+         389]
+
+open_cc_dids = [11,
+ 14,
+ 15,
+ 16,
+ 18,
+ 22,
+ 23,
+ 29,
+ 31,
+ 37,
+ 50,
+ 54,
+ 188,
+ 458,
+ 469,
+ 1049,
+ 1050,
+ 1063,
+ 1068,
+ 1510,
+ 1494,
+ 1480,
+ 1462,
+ 1464,
+ 6332,
+ 23381,
+ 40966,
+ 40982,
+ 40994,
+ 40975]
+# Filtered by N_samples < 2000, N feats < 100, N classes < 10
+
+open_cc_valid_dids = [13,25,35,40,41,43,48,49,51,53,55,56,59,61,187,285,329,333,334,335,336,337,338,377,446,450,451,452,460,463,464,466,470,475,481,679,694,717,721,724,733,738,745,747,748,750,753,756,757,764,765,767,774,778,786,788,795,796,798,801,802,810,811,814,820,825,826,827,831,839,840,841,844,852,853,854,860,880,886,895,900,906,907,908,909,915,925,930,931,934,939,940,941,949,966,968,984,987,996,1048,1054,1071,1073,1100,1115,1412,1442,1443,1444,1446,1447,1448,1451,1453,1488,1490,1495,1498,1499,1506,1508,1511,1512,1520,1523,4153,23499,40496,40646,40663,40669,40680,40682,40686,40690,40693,40705,40706,40710,40711,40981,41430,41538,41919,41976,42172,42261,42544,42585,42638]

TabPFN/datasets/utils.py

@@ -0,0 +1,8 @@
+def normalize_data(eval_xs):
+    mean = eval_xs.mean(0)
+    std = eval_xs.std(0) + .000001
+    eval_xs = (eval_xs - mean) / std
+
+    return eval_xs
+
+

TabPFN/decoders.py

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

TabPFN/differentiable_pfn_evaluation.py

@@ -0,0 +1,345 @@
+import os
+import torch
+import numpy as np
+import time
+import pickle
+from scripts import  tabular_metrics
+from scripts.tabular_metrics import calculate_score_per_method
+from scripts.tabular_evaluation import evaluate
+from priors.differentiable_prior import draw_random_style
+from tqdm import tqdm
+import random
+from scripts.transformer_prediction_interface import get_params_from_config, load_model_workflow
+
+"""
+===============================
+PUBLIC FUNCTIONS FOR EVALUATION
+===============================
+"""
+
+
+def eval_model_range(i_range, *args, **kwargs):
+    for i in i_range:
+        eval_model(i, *args, **kwargs)
+
+
+
+def eval_model(i, e, valid_datasets, test_datasets, train_datasets, eval_positions_valid, eval_positions_test,
+               bptt_valid,
+               bptt_test, add_name, base_path, device='cpu', eval_addition='', **extra_tuning_args):
+    """
+    Differentiable model evaliation workflow. Evaluates and saves results to disk.
+
+    :param i:
+    :param e:
+    :param valid_datasets:
+    :param test_datasets:
+    :param train_datasets:
+    :param eval_positions_valid:
+    :param eval_positions_test:
+    :param bptt_valid:
+    :param bptt_test:
+    :param add_name:
+    :param base_path:
+    :param device:
+    :param eval_addition:
+    :param extra_tuning_args:
+    :return:
+    """
+    model, c, results_file = load_model_workflow(i, e, add_name, base_path, device, eval_addition)
+    params = {'bptt': bptt_valid
+        , 'bptt_final': bptt_test
+        , 'eval_positions': eval_positions_valid
+        , 'eval_positions_test': eval_positions_test
+        , 'valid_datasets': valid_datasets
+        , 'test_datasets': test_datasets
+        , 'train_datasets': train_datasets
+        , 'verbose': True
+        , 'device': device
+              }
+
+    params.update(get_params_from_config(c))
+
+    start = time.time()
+    metrics, metrics_valid, style, temperature, optimization_route = evaluate_differentiable_model(model, **params,
+                                                                                                   **extra_tuning_args)
+    print('Evaluation time: ', time.time() - start)
+
+    print(results_file)
+    r = [c.copy(), metrics, metrics_valid, style.to('cpu'), temperature.to('cpu'), optimization_route]
+    with open(results_file, 'wb') as output:
+        del r[0]['num_features_used']
+        del r[0]['categorical_features_sampler']
+        pickle.dump(r, output)
+
+    _, _, _, style, temperature, _ = r
+
+    return r, model
+
+"""
+===============================
+INTERNAL HELPER FUNCTIONS
+===============================
+"""
+
+def evaluate_differentiable_model(model
+                                  , valid_datasets
+                                  , test_datasets
+                                  , train_datasets
+                                  , N_draws=100
+                                  , N_grad_steps=10
+                                  , eval_positions=None
+                                  , eval_positions_test=None
+                                  , bptt=100
+                                  , bptt_final=200
+                                  , style=None
+                                  , n_parallel_configurations=1
+                                  , device='cpu'
+                                  , selection_metric='auc'
+                                  , final_splits=[1, 2, 3, 4, 5]
+                                  , N_ensemble_configurations_list=[1, 5, 10, 20, 50, 100]
+                                  , **kwargs):
+    """
+    Evaluation function for diffable model evaluation. Returns a list of results.
+
+    :param model:
+    :param valid_datasets:
+    :param test_datasets:
+    :param train_datasets:
+    :param N_draws:
+    :param N_grad_steps:
+    :param eval_positions:
+    :param eval_positions_test:
+    :param bptt:
+    :param bptt_final:
+    :param style:
+    :param n_parallel_configurations:
+    :param device:
+    :param selection_metric:
+    :param final_splits:
+    :param N_ensemble_configurations_list:
+    :param kwargs:
+    :return:
+    """
+    torch.manual_seed(0)
+    np.random.seed(0)
+    random.seed(0)
+
+    diffable_metric = tabular_metrics.cross_entropy
+    evaluation_metric = tabular_metrics.auc_metric
+    if selection_metric in ('auc', 'roc'):
+        selection_metric_min_max = 'max'
+        selection_metric = tabular_metrics.auc_metric
+        evaluation_metric = selection_metric
+    elif selection_metric in ('ce', 'selection_metric'):
+        selection_metric_min_max = 'min'
+        selection_metric = tabular_metrics.cross_entropy
+        evaluation_metric = selection_metric
+
+    print('Diffable metric', diffable_metric, ' Selection metric', selection_metric, ' Evaluation metric',
+          evaluation_metric)
+    print('N PARALLEL CONFIGURATIONS', n_parallel_configurations)
+    print('eval_positions', eval_positions)
+
+    def evaluate_valid(style, softmax_temperature, results, results_tracked):
+        result_valid = eval_step(valid_datasets, style, softmax_temperature=softmax_temperature,
+                                 return_tensor=False, inference_mode=True, selection_metric=selection_metric,
+                                 evaluation_metric=evaluation_metric, eval_positions=eval_positions, bptt=bptt, model=model[2])
+        result_valid = [float(result_valid[f'mean_select_at_{pos}']) for pos in eval_positions]
+        results += [result_valid]
+        results_tracked += [np.nanmean(result_valid)]
+
+    model[2].to(device)
+    model[2].eval()
+
+    results_on_valid, results_on_valid_tracked = [], []
+    best_style, best_softmax_temperature = style, torch.cat(
+        [torch.tensor([0.0]).to(device) for n in range(0, n_parallel_configurations)], 0)
+    optimization_routes = []
+
+    best_style = torch.cat([draw_random_style(model[3], device).detach() for n in range(0, n_parallel_configurations)],
+                      0)
+    best_softmax_temperature = torch.cat([torch.tensor([0.0]).to(device) for n in range(0, n_parallel_configurations)],
+                                    0)
+
+
+    for _ in tqdm(range(0, N_draws), desc='Iterate over Optimization initializations'): # Evaluates N hparam draws
+        style = torch.cat([draw_random_style(model[3], device).detach() for n in range(0, n_parallel_configurations)],
+                          0)
+        softmax_temperature = torch.cat([torch.tensor([0.0]).to(device) for n in range(0, n_parallel_configurations)],
+                                        0)
+
+        evaluate_valid(style, softmax_temperature, results_on_valid, results_on_valid_tracked)
+
+        print(f'Draw --> Valid Selection metric: {results_on_valid[-1]}')
+
+        if N_grad_steps > 0:
+            gradient_optimize_result = gradient_optimize_style(model, style, N_grad_steps
+                                                               , softmax_temperature=softmax_temperature
+                                                               , model=model[2]
+                                                               , train_datasets=train_datasets
+                                                               , valid_datasets=valid_datasets
+                                                               , selection_metric_min_max=selection_metric_min_max
+                                                               , **kwargs)
+            optimization_routes += [gradient_optimize_result['optimization_route']]
+
+            evaluate_valid(gradient_optimize_result['best_style']
+                                          , gradient_optimize_result['best_temperature']
+                                          , results_on_valid, results_on_valid_tracked)
+
+            print(f'After diff --> Valid Selection metric: {results_on_valid[-1]}')
+
+        if selection_metric_min_max == 'min':
+            is_best = (results_on_valid_tracked[-1] <= min(results_on_valid_tracked))
+        else:
+            is_best = (results_on_valid_tracked[-1] >= max(results_on_valid_tracked))
+
+        if is_best or best_style is None:
+            best_style = gradient_optimize_result['best_style'].clone()
+            best_softmax_temperature = gradient_optimize_result['best_temperature'].clone()
+    torch.cuda.empty_cache()
+
+    def final_evaluation():
+        print('Running eval dataset with final params (no gradients)..')
+        print(best_style, best_softmax_temperature)
+        result_test = []
+        for N_ensemble_configurations in N_ensemble_configurations_list:
+            print(f'Running with {N_ensemble_configurations} ensemble_configurations')
+            kwargs['N_ensemble_configurations'] = N_ensemble_configurations
+            splits = []
+            for split in final_splits:
+                splits += [eval_step(test_datasets, best_style, softmax_temperature=best_softmax_temperature
+                                     , return_tensor=False, eval_positions=eval_positions_test,
+                                     bptt=bptt_final, inference_mode=True, split_number=split, model=model[2]
+                                     , selection_metric=selection_metric, evaluation_metric=evaluation_metric)]
+            result_test += [splits]
+
+        print('Running valid dataset with final params (no gradients)..')
+        result_valid = eval_step(valid_datasets, best_style, softmax_temperature=best_softmax_temperature
+                                 , return_tensor=False, eval_positions=eval_positions_test,
+                                 bptt=bptt_final, inference_mode=True, model=model[2]
+                                 , selection_metric=selection_metric, evaluation_metric=evaluation_metric)
+
+        return result_test, result_valid
+
+    result_test, result_valid = final_evaluation()
+
+    return result_test, result_valid, best_style, best_softmax_temperature, optimization_routes
+
+
+def eval_step(ds, used_style, selection_metric, evaluation_metric, eval_positions, return_tensor=True, **kwargs):
+    def step():
+        return evaluate(datasets=ds,
+                        method='transformer'
+                        , overwrite=True
+                        , style=used_style
+                        , eval_positions=eval_positions
+                        , metric_used=selection_metric
+                        , save=False
+                        , path_interfix=None
+                        , base_path=None
+                        , verbose=True
+                        , **kwargs)
+
+    if return_tensor:
+        r = step()
+    else:
+        with torch.no_grad():
+            r = step()
+
+    calculate_score_per_method(selection_metric, 'select', r, ds, eval_positions, aggregator='mean')
+    calculate_score_per_method(evaluation_metric, 'eval', r, ds, eval_positions, aggregator='mean')
+
+    return r
+
+
+def gradient_optimize_style(model, init_style, steps, softmax_temperature, train_datasets, valid_datasets, learning_rate=0.03, optimize_all=False,
+                            limit_style=True, N_datasets_sampled=90, optimize_softmax_temperature=True, selection_metric_min_max='max', **kwargs):
+    """
+    Uses gradient based methods to optimize 'style' on the 'train_datasets' and uses stopping with 'valid_datasets'.
+
+    :param model:
+    :param init_style:
+    :param steps:
+    :param learning_rate:
+    :param softmax_temperature:
+    :param train_datasets:
+    :param valid_datasets:
+    :param optimize_all:
+    :param limit_style:
+    :param N_datasets_sampled:
+    :param optimize_softmax_temperature:
+    :param selection_metric_min_max:
+    :param kwargs:
+    :return:
+    """
+    grad_style = torch.nn.Parameter(init_style.detach(), requires_grad=True)
+
+    best_style, best_temperature, best_selection_metric, best_diffable_metric = grad_style.detach(), softmax_temperature.detach(), None, None
+    softmax_temperature = torch.nn.Parameter(softmax_temperature.detach(), requires_grad=optimize_softmax_temperature)
+    variables_to_optimize = model[2].parameters() if optimize_all else [grad_style, softmax_temperature]
+    optimizer = torch.optim.Adam(variables_to_optimize, lr=learning_rate)
+
+    optimization_route_selection, optimization_route_diffable = [], []
+    optimization_route_selection_valid, optimization_route_diffable_valid = [], []
+
+    def eval_opt(ds, return_tensor=True, inference_mode=False):
+        result = eval_step(ds, grad_style, softmax_temperature=softmax_temperature, return_tensor=return_tensor
+                           , inference_mode=inference_mode, model=model[2], **kwargs)
+
+        diffable_metric = result['mean_metric']
+        selection_metric = result['mean_select']
+
+        return diffable_metric, selection_metric
+
+    def eval_all_datasets(datasets, propagate=True):
+        selection_metrics_this_step, diffable_metrics_this_step = [], []
+        for ds in datasets:
+            diffable_metric_train, selection_metric_train = eval_opt([ds], inference_mode=(not propagate))
+            if not torch.isnan(diffable_metric_train).any():
+                if propagate and diffable_metric_train.requires_grad == True:
+                    diffable_metric_train.backward()
+                selection_metrics_this_step += [selection_metric_train]
+                diffable_metrics_this_step += [float(diffable_metric_train.detach().cpu().numpy())]
+        diffable_metric_train = np.nanmean(diffable_metrics_this_step)
+        selection_metric_train = np.nanmean(selection_metrics_this_step)
+
+        return diffable_metric_train, selection_metric_train
+
+    for t in tqdm(range(steps), desc='Iterate over Optimization steps'):
+        optimizer.zero_grad()
+
+        # Select subset of datasets
+        random.seed(t)
+        train_datasets_ = random.sample(train_datasets, N_datasets_sampled)
+
+        # Get score on train
+        diffable_metric_train, selection_metric_train = eval_all_datasets(train_datasets_, propagate=True)
+        optimization_route_selection += [float(selection_metric_train)]
+        optimization_route_diffable += [float(diffable_metric_train)]
+
+        # Get score on valid
+        diffable_metric_valid, selection_metric_valid = eval_all_datasets(valid_datasets, propagate=False)
+        optimization_route_selection_valid += [float(selection_metric_valid)]
+        optimization_route_diffable_valid += [float(diffable_metric_valid)]
+
+        is_best = (selection_metric_min_max == 'min' and best_selection_metric > selection_metric_valid)
+        is_best = is_best or (selection_metric_min_max == 'max' and best_selection_metric < selection_metric_valid)
+        if (best_selection_metric is None) or (not np.isnan(selection_metric_valid) and is_best):
+            print('New best', best_selection_metric, selection_metric_valid)
+            best_style = grad_style.detach().clone()
+            best_temperature = softmax_temperature.detach().clone()
+            best_selection_metric, best_diffable_metric = selection_metric_valid, diffable_metric_valid
+
+        optimizer.step()
+
+        if limit_style:
+            grad_style = grad_style.detach().clamp(-1.74, 1.74)
+
+        print(f'Valid: Diffable metric={diffable_metric_valid} Selection metric={selection_metric_valid};' +
+            f'Train: Diffable metric={diffable_metric_train} Selection metric={selection_metric_train}')
+
+    print(f'Return best:{best_style} {best_selection_metric}')
+    return {'best_style': best_style, 'best_temperature': best_temperature
+            , 'optimization_route': {'select': optimization_route_selection, 'loss': optimization_route_diffable,
+               'test_select': optimization_route_selection_valid, 'test_loss': optimization_route_diffable_valid}}

TabPFN/encoders.py

@@ -0,0 +1,225 @@
+import math
+
+import torch
+import torch.nn as nn
+from utils import normalize_data
+import torch.nn.functional as F
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
+
+
+class StyleEncoder(nn.Module):
+    def __init__(self, em_size, hyperparameter_definitions):
+        super().__init__()
+        # self.embeddings = {}
+        self.em_size = em_size
+        # self.hyperparameter_definitions = {}
+        # for hp in hyperparameter_definitions:
+        #     self.embeddings[hp] = nn.Linear(1, self.em_size)
+        # self.embeddings = nn.ModuleDict(self.embeddings)
+        self.embedding = nn.Linear(hyperparameter_definitions.shape[0], self.em_size)
+
+    def forward(self, hyperparameters):  # T x B x num_features
+        # Make faster by using matrices
+        # sampled_embeddings = [torch.stack([
+        #     self.embeddings[hp](torch.tensor([batch[hp]], device=self.embeddings[hp].weight.device, dtype=torch.float))
+        #     for hp in batch
+        # ], -1).sum(-1) for batch in hyperparameters]
+        # return torch.stack(sampled_embeddings, 0)
+        return self.embedding(hyperparameters)
+
+
+class _PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout=0.):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        self.d_model = d_model
+        self.device_test_tensor = nn.Parameter(torch.tensor(1.))
+
+    def forward(self, x):# T x B x num_features
+        assert self.d_model % x.shape[-1]*2 == 0
+        d_per_feature = self.d_model // x.shape[-1]
+        pe = torch.zeros(*x.shape, d_per_feature, device=self.device_test_tensor.device)
+        #position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        interval_size = 10
+        div_term = (1./interval_size) * 2*math.pi*torch.exp(torch.arange(0, d_per_feature, 2, device=self.device_test_tensor.device).float()*math.log(math.sqrt(2)))
+        #print(div_term/2/math.pi)
+        pe[..., 0::2] = torch.sin(x.unsqueeze(-1) * div_term)
+        pe[..., 1::2] = torch.cos(x.unsqueeze(-1) * div_term)
+        return self.dropout(pe).view(x.shape[0],x.shape[1],self.d_model)
+
+
+Positional = lambda _, emsize: _PositionalEncoding(d_model=emsize)
+
+class EmbeddingEncoder(nn.Module):
+    def __init__(self, num_features, em_size, num_embs=100):
+        super().__init__()
+        self.num_embs = num_embs
+        self.embeddings = nn.Embedding(num_embs * num_features, em_size, max_norm=True)
+        self.init_weights(.1)
+        self.min_max = (-2,+2)
+
+    @property
+    def width(self):
+        return self.min_max[1] - self.min_max[0]
+
+    def init_weights(self, initrange):
+        self.embeddings.weight.data.uniform_(-initrange, initrange)
+
+    def discretize(self, x):
+        split_size = self.width / self.num_embs
+        return (x - self.min_max[0] // split_size).int().clamp(0, self.num_embs - 1)
+
+    def forward(self, x):  # T x B x num_features
+        x_idxs = self.discretize(x)
+        x_idxs += torch.arange(x.shape[-1], device=x.device).view(1, 1, -1) * self.num_embs
+        # print(x_idxs,self.embeddings.weight.shape)
+        return self.embeddings(x_idxs).mean(-2)
+
+
+class Normalize(nn.Module):
+    def __init__(self, mean, std):
+        super().__init__()
+        self.mean = mean
+        self.std = std
+
+    def forward(self, x):
+        return (x-self.mean)/self.std
+
+
+def get_normalized_uniform_encoder(encoder_creator):
+    """
+    This can be used to wrap an encoder that is fed uniform samples in [0,1] and normalizes these to 0 mean and 1 std.
+    For example, it can be used as `encoder_creator = get_normalized_uniform_encoder(encoders.Linear)`, now this can
+    be initialized with `encoder_creator(feature_dim, in_dim)`.
+    :param encoder:
+    :return:
+    """
+    return lambda in_dim, out_dim: nn.Sequential(Normalize(.5, math.sqrt(1/12)), encoder_creator(in_dim, out_dim))
+
+
+Linear = nn.Linear
+MLP = lambda num_features, emsize: nn.Sequential(nn.Linear(num_features+1,emsize*2),
+                                                 nn.ReLU(),
+                                                 nn.Linear(emsize*2,emsize))
+
+class NanHandlingEncoder(nn.Module):
+    def __init__(self, num_features, emsize, keep_nans=True):
+        super().__init__()
+        self.num_features = 2 * num_features if keep_nans else num_features
+        self.emsize = emsize
+        self.keep_nans = keep_nans
+        self.layer = nn.Linear(self.num_features, self.emsize)
+
+    def forward(self, x):
+        if self.keep_nans:
+            x = torch.cat([torch.nan_to_num(x, nan=0.0), normalize_data(torch.isnan(x) * -1
+                                                          + torch.logical_and(torch.isinf(x), torch.sign(x) == 1) * 1
+                                                          + torch.logical_and(torch.isinf(x), torch.sign(x) == -1) * 2
+                                                          )], -1)
+        else:
+            x = torch.nan_to_num(x, nan=0.0)
+        return self.layer(x)
+
+class Linear(nn.Linear):
+    def __init__(self, num_features, emsize):
+        super().__init__(num_features, emsize)
+        self.num_features = num_features
+        self.emsize = emsize
+
+    def forward(self, x):
+        x = torch.nan_to_num(x, nan=0.0)
+        return super().forward(x)
+
+class SequenceSpanningEncoder(nn.Module):
+    # Regular Encoder transforms Seq_len, B, S -> Seq_len, B, E attending only to last dimension
+    # This Encoder accesses the Seq_Len dimension additionally
+
+    # Why would we want this? We can learn normalization and embedding of features
+    #    , this might be more important for e.g. categorical, ordinal feats, nan detection
+    # However maybe this can be easily learned through transformer as well?
+    # A problem is to make this work across any sequence length and be independent of ordering
+
+    # We could use average and maximum pooling and use those with a linear layer
+
+
+    # Another idea !! Similar to this we would like to encode features so that their number is variable
+    # We would like to embed features, also using knowledge of the features in the entire sequence
+
+    # We could use convolution or another transformer
+    # Convolution:
+
+    # Transformer/Conv across sequence dimension that encodes and normalizes features
+    #    -> Transformer across feature dimension that encodes features to a constant size
+
+    # Conv with flexible features but no sequence info: S,B,F -(reshape)-> S*B,1,F
+    #   -(Conv1d)-> S*B,N,F -(AvgPool,MaxPool)-> S*B,N,1 -> S,B,N
+    # This probably won't work since it's missing a way to recognize which feature is encoded
+
+    # Transformer with flexible features: S,B,F -> F,B*S,1 -> F2,B*S,1 -> S,B,F2
+
+    def __init__(self, num_features, em_size):
+        super().__init__()
+
+        raise NotImplementedError()
+        # Seq_len, B, S -> Seq_len, B, E
+        #
+        self.convs = torch.nn.ModuleList([nn.Conv1d(64 if i else 1, 64, 3) for i in range(5)])
+        # self.linear = nn.Linear(64, emsize)
+
+class TransformerBasedFeatureEncoder(nn.Module):
+    def __init__(self, num_features, emsize):
+        super().__init__()
+
+        hidden_emsize = emsize
+        encoder = Linear(1, hidden_emsize)
+        n_out = emsize
+        nhid = 2*emsize
+        dropout =0.0
+        nhead=4
+        nlayers=4
+        model = nn.Transformer(nhead=nhead, num_encoder_layers=4, num_decoder_layers=4, d_model=1)
+
+    def forward(self, *input):
+        # S,B,F -> F,S*B,1 -> F2,S*B,1 -> S,B,F2
+        input = input.transpose()
+        self.model(input)
+
+class Conv(nn.Module):
+    def __init__(self, input_size, emsize):
+        super().__init__()
+        self.convs = torch.nn.ModuleList([nn.Conv2d(64 if i else 1, 64, 3) for i in range(5)])
+        self.linear = nn.Linear(64,emsize)
+
+
+    def forward(self, x):
+        size = math.isqrt(x.shape[-1])
+        assert size*size == x.shape[-1]
+        x = x.reshape(*x.shape[:-1], 1, size, size)
+        for conv in self.convs:
+            if x.shape[-1] < 4:
+                break
+            x = conv(x)
+            x.relu_()
+        x = nn.AdaptiveAvgPool2d((1,1))(x).squeeze(-1).squeeze(-1)
+        return self.linear(x)
+
+
+
+
+class CanEmb(nn.Embedding):
+    def __init__(self, num_features, num_embeddings: int, embedding_dim: int, *args, **kwargs):
+        assert embedding_dim % num_features == 0
+        embedding_dim = embedding_dim // num_features
+        super().__init__(num_embeddings, embedding_dim, *args, **kwargs)
+
+    def forward(self, x):
+        lx = x.long()
+        assert (lx == x).all(), "CanEmb only works with tensors of whole numbers"
+        x = super().forward(lx)
+        return x.view(*x.shape[:-2], -1)
+
+def get_Canonical(num_classes):
+    return lambda num_features, emsize: CanEmb(num_features, num_classes, emsize)
+
+def get_Embedding(num_embs_per_feature=100):
+    return lambda num_features, emsize: EmbeddingEncoder(num_features, emsize, num_embs=num_embs_per_feature)

TabPFN/initializers.py

@@ -0,0 +1,9 @@
+import torch
+from torch import nn
+
+def get_NormalInitializer(std):
+    def initializer(m):
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, 0, std)
+            nn.init.normal_(m.bias, 0, std)
+    return initializer

TabPFN/layer.py

@@ -0,0 +1,125 @@
+from functools import partial
+
+from torch import nn
+from torch.nn.modules.transformer import *
+from torch.nn.modules.transformer import _get_activation_fn
+
+from torch.utils.checkpoint import checkpoint
+
+
+class TransformerEncoderLayer(Module):
+    r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
+    This standard encoder layer is based on the paper "Attention Is All You Need".
+    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
+    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
+    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
+    in a different way during application.
+
+    Args:
+        d_model: the number of expected features in the input (required).
+        nhead: the number of heads in the multiheadattention models (required).
+        dim_feedforward: the dimension of the feedforward network model (default=2048).
+        dropout: the dropout value (default=0.1).
+        activation: the activation function of intermediate layer, relu or gelu (default=relu).
+        layer_norm_eps: the eps value in layer normalization components (default=1e-5).
+        batch_first: If ``True``, then the input and output tensors are provided
+            as (batch, seq, feature). Default: ``False``.
+
+    Examples::
+        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = encoder_layer(src)
+
+    Alternatively, when ``batch_first`` is ``True``:
+        >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)
+        >>> src = torch.rand(32, 10, 512)
+        >>> out = encoder_layer(src)
+    """
+    __constants__ = ['batch_first']
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu",
+                 layer_norm_eps=1e-5, batch_first=False, pre_norm=False,
+                 device=None, dtype=None, recompute_attn=False) -> None:
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        super().__init__()
+        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
+                                            **factory_kwargs)
+        # Implementation of Feedforward model
+        self.linear1 = Linear(d_model, dim_feedforward, **factory_kwargs)
+        self.dropout = Dropout(dropout)
+        self.linear2 = Linear(dim_feedforward, d_model, **factory_kwargs)
+
+        self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
+        self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
+        self.dropout1 = Dropout(dropout)
+        self.dropout2 = Dropout(dropout)
+        self.pre_norm = pre_norm
+        self.recompute_attn = recompute_attn
+
+        self.activation = _get_activation_fn(activation)
+
+    def __setstate__(self, state):
+        if 'activation' not in state:
+            state['activation'] = F.relu
+        super().__setstate__(state)
+
+    def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            src: the sequence to the encoder layer (required).
+            src_mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        if self.pre_norm:
+            src_ = self.norm1(src)
+        else:
+            src_ = src
+        if isinstance(src_mask, tuple):
+            # global attention setup
+            assert not self.self_attn.batch_first
+            assert src_key_padding_mask is None
+
+            global_src_mask, trainset_src_mask, valset_src_mask = src_mask
+
+            num_global_tokens = global_src_mask.shape[0]
+            num_train_tokens = trainset_src_mask.shape[0]
+
+            global_tokens_src = src_[:num_global_tokens]
+            train_tokens_src = src_[num_global_tokens:num_global_tokens+num_train_tokens]
+            global_and_train_tokens_src = src_[:num_global_tokens+num_train_tokens]
+            eval_tokens_src = src_[num_global_tokens+num_train_tokens:]
+
+
+            attn = partial(checkpoint, self.self_attn) if self.recompute_attn else self.self_attn
+
+            global_tokens_src2 = attn(global_tokens_src, global_and_train_tokens_src, global_and_train_tokens_src, None, True, global_src_mask)[0]
+            train_tokens_src2 = attn(train_tokens_src, global_tokens_src, global_tokens_src, None, True, trainset_src_mask)[0]
+            eval_tokens_src2 = attn(eval_tokens_src, src_, src_,
+                                    None, True, valset_src_mask)[0]
+
+            src2 = torch.cat([global_tokens_src2, train_tokens_src2, eval_tokens_src2], dim=0)
+
+        else:
+            if self.recompute_attn:
+                src2 = checkpoint(self.self_attn, src_, src_, src_, src_key_padding_mask, True, src_mask)[0]
+            else:
+                src2 = self.self_attn(src_, src_, src_, attn_mask=src_mask,
+                                      key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        if not self.pre_norm:
+            src = self.norm1(src)
+
+        if self.pre_norm:
+            src_ = self.norm2(src)
+        else:
+            src_ = src
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src_))))
+        src = src + self.dropout2(src2)
+
+        if not self.pre_norm:
+            src = self.norm2(src)
+        return src

TabPFN/losses.py

@@ -0,0 +1,41 @@
+import torch
+from torch import nn
+
+class CrossEntropyForMulticlassLoss(torch.nn.CrossEntropyLoss):
+    # This loss applies cross entropy after reducing the number of prediction
+    # dimensions to the number of classes in the target
+
+    # TODO: loss.item() doesn't work so the displayed losses are Nans
+    def __init__(self, num_classes, weight=None, size_average=None, ignore_index: int = -100,
+                 reduce=None, reduction: str = 'mean', label_smoothing: float = 0.0) -> None:
+        super().__init__(size_average=size_average, reduce=reduce, reduction=reduction, ignore_index=ignore_index)
+        self.num_classes = num_classes
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        loss = torch.zeros_like(input[:, :, 0])
+        for b in range(target.shape[1]):
+            l = super().forward(input[:, b, 0:len(torch.unique(target[:, b]))], target[:, b])
+            loss[:, b] += l
+        return loss.flatten()
+
+def JointBCELossWithLogits(output, target):
+    # output shape: (S, B, NS) with NS = Number of sequences
+    # target shape: (S, B, SL)
+    # Loss = -log(mean_NS(prod_SL(p(target_SL, output_NS))))
+    # Here at the moment NS = SL
+    output = output.unsqueeze(-1).repeat(1, 1, 1, target.shape[-1]) # (S, B, NS, SL)
+    output = output.permute(2, 0, 1, 3) # (NS, S, B, SL)
+    print(target.shape, output.shape)
+    loss = (target * torch.sigmoid(output)) + ((1-target) * (1-torch.sigmoid(output)))
+    loss = loss.prod(-1)
+    loss = loss.mean(0)
+    loss = -torch.log(loss)
+    loss = loss.mean()
+    return loss
+
+class ScaledSoftmaxCE(nn.Module):
+    def forward(self, x, label):
+        logits = x[..., :-10]
+        temp_scales = x[..., -10:]
+
+        logprobs = logits.softmax(-1)

TabPFN/model_builder.py

@@ -0,0 +1,273 @@
+from train import train, Losses
+import priors
+import encoders
+
+from collections import defaultdict
+
+from priors.utils import trunc_norm_sampler_f, gamma_sampler_f
+from utils import get_uniform_single_eval_pos_sampler
+import torch
+import math
+
+def save_model(model, path, filename, config_sample):
+    config_sample = {**config_sample}
+
+    def make_serializable(config_sample):
+        if isinstance(config_sample, dict):
+            config_sample = {k: make_serializable(config_sample[k]) for k in config_sample}
+        if isinstance(config_sample, list):
+            config_sample = [make_serializable(v) for v in config_sample]
+        if callable(config_sample):
+            config_sample = str(config_sample)
+        return config_sample
+
+    #if 'num_features_used' in config_sample:
+    #    del config_sample['num_features_used']
+
+    #config_sample['num_classes_as_str'] = str(config_sample['num_classes'])
+    #del config_sample['num_classes']
+
+    config_sample = make_serializable(config_sample)
+
+    torch.save((model.state_dict(), None, config_sample), os.path.join(path, filename))
+
+
+import subprocess as sp
+import os
+
+def get_gpu_memory():
+    command = "nvidia-smi"
+    memory_free_info = sp.check_output(command.split()).decode('ascii')
+    return memory_free_info
+
+
+def load_model(path, filename, device, eval_positions, verbose):
+    # TODO: This function only restores evaluation functionality but training canät be continued. It is also not flexible.
+
+    model_state, optimizer_state, config_sample = torch.load(
+        os.path.join(path, filename), map_location='cpu')
+    if ('differentiable_hyperparameters' in config_sample
+            and 'prior_mlp_activations' in config_sample['differentiable_hyperparameters']):
+        config_sample['differentiable_hyperparameters']['prior_mlp_activations']['choice_values_used'] = config_sample[
+                                                                                                         'differentiable_hyperparameters'][
+                                                                                                         'prior_mlp_activations'][
+                                                                                                         'choice_values']
+        config_sample['differentiable_hyperparameters']['prior_mlp_activations']['choice_values'] = [
+            torch.nn.Tanh for k in config_sample['differentiable_hyperparameters']['prior_mlp_activations']['choice_values']]
+
+    config_sample['categorical_features_sampler'] = lambda: lambda x: ([], [], [])
+    config_sample['num_features_used_in_training'] = config_sample['num_features_used']
+    config_sample['num_features_used'] = lambda: config_sample['num_features']
+    config_sample['num_classes_in_training'] = config_sample['num_classes']
+    config_sample['num_classes'] = 2
+    config_sample['batch_size_in_training'] = config_sample['batch_size']
+    config_sample['batch_size'] = 1
+    config_sample['bptt_in_training'] = config_sample['bptt']
+    config_sample['bptt'] = 10
+    config_sample['bptt_extra_samples_in_training'] = config_sample['bptt_extra_samples']
+    config_sample['bptt_extra_samples'] = None
+
+    #print('Memory', str(get_gpu_memory()))
+
+    model = get_model(config_sample, device=device, should_train=False, verbose=verbose)
+    module_prefix = 'module.'
+    model_state = {k.replace(module_prefix, ''): v for k, v in model_state.items()}
+    model[2].load_state_dict(model_state)
+    model[2].to(device)
+
+    return model, config_sample
+
+def fix_loaded_config_sample(loaded_config_sample, config):
+    def copy_to_sample(*k):
+        t,s = loaded_config_sample, config
+        for k_ in k[:-1]:
+            t = t[k_]
+            s = s[k_]
+        t[k[-1]] = s[k[-1]]
+    copy_to_sample('num_features_used')
+    copy_to_sample('num_classes')
+    copy_to_sample('differentiable_hyperparameters','prior_mlp_activations','choice_values')
+
+def load_config_sample(path, template_config):
+    model_state, optimizer_state, loaded_config_sample = torch.load(path, map_location='cpu')
+    fix_loaded_config_sample(loaded_config_sample, template_config)
+    return loaded_config_sample
+
+def get_default_spec(test_datasets, valid_datasets):
+    bptt = 10000
+    eval_positions = [1000, 2000, 3000, 4000, 5000] # list(2 ** np.array([4, 5, 6, 7, 8, 9, 10, 11, 12]))
+    max_features = max([X.shape[1] for (_, X, _, _, _, _) in test_datasets] + [X.shape[1] for (_, X, _, _, _, _) in valid_datasets])
+    max_splits = 5
+
+    return bptt, eval_positions, max_features, max_splits
+
+def get_mlp_prior_hyperparameters(config):
+    config = {hp: (list(config[hp].values())[0]) if type(config[hp]) is dict else config[hp] for hp in config}
+
+    if "prior_sigma_gamma_k" in config:
+        sigma_sampler = gamma_sampler_f(config["prior_sigma_gamma_k"], config["prior_sigma_gamma_theta"])
+        config['init_std'] = sigma_sampler
+    if "prior_noise_std_gamma_k" in config:
+        noise_std_sampler = gamma_sampler_f(config["prior_noise_std_gamma_k"], config["prior_noise_std_gamma_theta"])
+        config['noise_std'] = noise_std_sampler
+
+    return config
+
+
+def get_gp_mix_prior_hyperparameters(config):
+    return {'lengthscale_concentration': config["prior_lengthscale_concentration"],
+            'nu': config["prior_nu"],
+            'outputscale_concentration': config["prior_outputscale_concentration"],
+            'categorical_data': config["prior_y_minmax_norm"],
+            'y_minmax_norm': config["prior_lengthscale_concentration"],
+            'noise_concentration': config["prior_noise_concentration"],
+            'noise_rate': config["prior_noise_rate"]}
+
+def get_gp_prior_hyperparameters(config):
+    return {hp: (list(config[hp].values())[0]) if type(config[hp]) is dict else config[hp] for hp in config}
+
+
+def get_meta_gp_prior_hyperparameters(config):
+    config = {hp: (list(config[hp].values())[0]) if type(config[hp]) is dict else config[hp] for hp in config}
+
+    if "outputscale_mean" in config:
+        outputscale_sampler = trunc_norm_sampler_f(config["outputscale_mean"]
+                                                   , config["outputscale_mean"] * config["outputscale_std_f"])
+        config['outputscale'] = outputscale_sampler
+    if "lengthscale_mean" in config:
+        lengthscale_sampler = trunc_norm_sampler_f(config["lengthscale_mean"],
+                                                   config["lengthscale_mean"] * config["lengthscale_std_f"])
+        config['lengthscale'] = lengthscale_sampler
+
+    return config
+
+
+def get_model(config, device, should_train=True, verbose=False, state_dict=None, epoch_callback=None):
+    extra_kwargs = {}
+    verbose_train, verbose_prior = verbose >= 1, verbose >= 2
+    config['verbose'] = verbose_prior
+
+    if 'aggregate_k_gradients' not in config or config['aggregate_k_gradients'] is None:
+        config['aggregate_k_gradients'] = math.ceil(config['batch_size'] * ((config['nlayers'] * config['emsize'] * config['bptt'] * config['bptt']) / 10824640000))
+
+    config['num_steps'] = math.ceil(config['num_steps'] * config['aggregate_k_gradients'])
+    config['batch_size'] = math.ceil(config['batch_size'] / config['aggregate_k_gradients'])
+    config['recompute_attn'] = config['recompute_attn'] if 'recompute_attn' in config else False
+
+    def make_get_batch(model_proto, **extra_kwargs):
+        extra_kwargs = defaultdict(lambda: None, **extra_kwargs)
+        return (lambda batch_size, seq_len, num_features, hyperparameters
+                , device, model_proto=model_proto, get_batch=extra_kwargs['get_batch']
+                       , prior_bag_priors=extra_kwargs['prior_bag_priors']: model_proto.get_batch(
+            batch_size=batch_size
+            , seq_len=seq_len
+            , device=device
+            , get_batch=get_batch
+            , hyperparameters=hyperparameters
+            , num_features=num_features))
+
+    if config['prior_type'] == 'prior_bag':
+        # Prior bag combines priors
+        get_batch_gp = make_get_batch(priors.fast_gp)
+        get_batch_mlp = make_get_batch(priors.mlp)
+        if 'flexible' in config and config['flexible']:
+            get_batch_gp = make_get_batch(priors.flexible_categorical, **{'get_batch': get_batch_gp})
+            get_batch_mlp = make_get_batch(priors.flexible_categorical, **{'get_batch': get_batch_mlp})
+        prior_bag_hyperparameters = {'prior_bag_get_batch': (get_batch_gp, get_batch_mlp)
+            , 'prior_bag_exp_weights_1': 2.0}
+        prior_hyperparameters = {**get_mlp_prior_hyperparameters(config), **get_gp_prior_hyperparameters(config)
+            , **prior_bag_hyperparameters}
+        model_proto = priors.prior_bag
+    else:
+        if config['prior_type'] == 'mlp':
+            prior_hyperparameters = get_mlp_prior_hyperparameters(config)
+            model_proto = priors.mlp
+        elif config['prior_type'] == 'gp':
+            prior_hyperparameters = get_gp_prior_hyperparameters(config)
+            model_proto = priors.fast_gp
+        elif config['prior_type'] == 'gp_mix':
+            prior_hyperparameters = get_gp_mix_prior_hyperparameters(config)
+            model_proto = priors.fast_gp_mix
+        else:
+            raise Exception()
+
+        if 'flexible' in config and config['flexible']:
+            get_batch_base = make_get_batch(model_proto)
+            extra_kwargs['get_batch'] = get_batch_base
+            model_proto = priors.flexible_categorical
+
+    use_style = False
+
+    if 'differentiable' in config and config['differentiable']:
+        get_batch_base = make_get_batch(model_proto, **extra_kwargs)
+        extra_kwargs = {'get_batch': get_batch_base, 'differentiable_hyperparameters': config['differentiable_hyperparameters']}
+        model_proto = priors.differentiable_prior
+        use_style = True
+    print(f"Using style prior: {use_style}")
+
+    if (('nan_prob_no_reason' in config and config['nan_prob_no_reason'] > 0.0) or
+        ('nan_prob_a_reason' in config and config['nan_prob_a_reason'] > 0.0) or
+        ('nan_prob_unknown_reason' in config and config['nan_prob_unknown_reason'] > 0.0)):
+        encoder = encoders.NanHandlingEncoder
+    else:
+        encoder = encoders.Linear
+
+    num_outputs = config['num_outputs'] if 'num_outputs' in config else 1
+    if config['max_num_classes'] == 2:
+        if 'joint_loss' in config and config['joint_loss']:
+            loss = JointBCELossWithLogits
+        else:
+            loss = Losses.bce
+    elif config['max_num_classes'] > 2:
+        loss = Losses.ce(torch.ones((config['max_num_classes'])))
+    else:
+        loss = BarDistribution(borders=get_bucket_limits(500, full_range=(-10, 10)))
+
+    aggregate_k_gradients = 1 if 'aggregate_k_gradients' not in config else config['aggregate_k_gradients']
+    check_is_compatible = False if 'multiclass_loss_type' not in config else (config['multiclass_loss_type'] == 'compatible')
+    config['multiclass_type'] = config['multiclass_type'] if 'multiclass_type' in config else 'rank'
+    config['mix_activations'] = config['mix_activations'] if 'mix_activations' in config else False
+
+    config['bptt_extra_samples'] = config['bptt_extra_samples'] if 'bptt_extra_samples' in config else None
+    config['eval_positions'] = [int(config['bptt'] * 0.95)] if config['bptt_extra_samples'] is None else [int(config['bptt'])]
+
+    epochs = 0 if not should_train else config['epochs']
+    model = train(model_proto.DataLoader
+                  , loss
+                  , encoder
+                  , style_encoder_generator = encoders.StyleEncoder if use_style else None
+                  , emsize=config['emsize']
+                  , nhead=config['nhead']
+                  , y_encoder_generator= encoders.get_Canonical(config['max_num_classes']) if config.get('canonical_y_encoder', False) else encoders.Linear
+                  , pos_encoder_generator=None
+                  , batch_size=config['batch_size']
+                  , nlayers=config['nlayers']
+                  , nhid=config['emsize'] * config['nhid_factor']
+                  , epochs=epochs
+                  , total_available_time_in_s=config.get('total_available_time_in_s', None)
+                  , warmup_epochs=20
+                  , bptt=config['bptt']
+                  , gpu_device=device
+                  , dropout=config['dropout']
+                  , steps_per_epoch=config['num_steps']
+                  , single_eval_pos_gen=get_uniform_single_eval_pos_sampler(config['bptt'])
+                  , load_weights_from_this_state_dict=state_dict
+                  , aggregate_k_gradients=aggregate_k_gradients
+                  , check_is_compatible=check_is_compatible
+                  , recompute_attn=config['recompute_attn']
+                  , epoch_callback=epoch_callback
+                  , bptt_extra_samples = config['bptt_extra_samples']
+                  , extra_prior_kwargs_dict={
+            'num_features': config['num_features']
+            , 'fuse_x_y': False
+            , 'hyperparameters': prior_hyperparameters
+            , 'num_outputs':num_outputs
+            , 'dynamic_batch_size': 1 if ('num_global_att_tokens' in config and config['num_global_att_tokens']) else 2
+            , **extra_kwargs
+        }
+                  , lr=config['lr']
+                  , verbose=verbose_train,
+                  weight_decay=config.get('weight_decay', 0.0),
+                  normalize_labels=True)
+
+    return model

TabPFN/models_diff/gp_ablation_model.cpkt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c7b0c8febc553cca3fdee265b5a1cd7567dbf83da855969940be4707a9218ffb
+size 69460013

TabPFN/models_diff/prior_diff_real_checkpoint_n_8x_lr0.0003_epoch_49.cpkt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dae97f45bd53d719fc2b23fac4ec55eab16d63892196d939b1bb1c3b408be242
+size 103616779

TabPFN/notebook_utils.py

@@ -0,0 +1,32 @@
+import os
+from pathlib import Path
+
+import io
+import torch
+import pickle
+
+def print_models(base_path, model_string):
+    print(model_string)
+
+    for i in range(80):
+        for e in range(50):
+            exists = Path(os.path.join(base_path, f'models_diff/prior_diff_real_checkpoint{model_string}_n_{i}_epoch_{e}.cpkt')).is_file()
+            if exists:
+                print(os.path.join(base_path, f'models_diff/prior_diff_real_checkpoint{model_string}_n_{i}_epoch_{e}.cpkt'))
+        print()
+
+class CustomUnpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if name == 'Manager':
+            from settings import Manager
+            return Manager
+        try:
+            return self.find_class_cpu(module, name)
+        except:
+            return None
+
+    def find_class_cpu(self, module, name):
+        if module == 'torch.storage' and name == '_load_from_bytes':
+            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
+        else:
+            return super().find_class(module, name)

TabPFN/positional_encodings.py

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

TabPFN/prior_tuning_result.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24d2189bbc836aeea888cf6c540f2c1b45b5351822931189e8bf10a0bc80a0b6
+size 18668851

TabPFN/priors/__init__.py

@@ -0,0 +1,4 @@
+from . import fast_gp, mlp, flexible_categorical, differentiable_prior, prior_bag
+
+
+

TabPFN/priors/differentiable_prior.py

@@ -0,0 +1,293 @@
+import torch
+from torch import nn
+import math
+
+from .utils import get_batch_to_dataloader
+from utils import default_device
+from .utils import order_by_y, normalize_by_used_features_f
+
+from .utils import trunc_norm_sampler_f, beta_sampler_f, gamma_sampler_f, uniform_sampler_f, zipf_sampler_f, scaled_beta_sampler_f, uniform_int_sampler_f
+
+
+def unpack_dict_of_tuples(d):
+    # Returns list of dicts where each dict i contains values of tuple position i
+    # {'a': (1,2), 'b': (3,4)} -> [{'a': 1, 'b': 3}, {'a': 2, 'b': 4}]
+    return [dict(zip(d.keys(), v)) for v in list(zip(*list(d.values())))]
+
+class DifferentiableHyperparameter(nn.Module):
+    ## We can sample this and get a hyperparameter value and a normalized hyperparameter indicator
+    def __init__(self, distribution, embedding_dim, device, **args):
+        super(DifferentiableHyperparameter, self).__init__()
+
+        self.distribution = distribution
+        self.embedding_dim = embedding_dim
+        self.device=device
+        for key in args:
+            setattr(self, key, args[key])
+
+        def get_sampler():
+            #if self.distribution == "beta":
+            #    return beta_sampler_f(self.a, self.b), 0, 1
+            #elif self.distribution == "gamma":
+            #    return gamma_sampler_f(self.a, self.b), 0, 1
+            #elif self.distribution == "beta_int":
+            #    return scaled_beta_sampler_f(self.a, self.b, self.scale, self.min), self.scale + self.min, self.min, self.a / (self.a + self.b)
+            if self.distribution == "uniform":
+                if not hasattr(self, 'sample'):
+                    return uniform_sampler_f(self.min, self.max), self.min, self.max, (self.max+self.min) / 2, math.sqrt(1/12*(self.max-self.min)*(self.max-self.min))
+                else:
+                    return lambda: self.sample, self.min, self.max, None, None
+            elif self.distribution == "uniform_int":
+                return uniform_int_sampler_f(self.min, self.max), self.min, self.max, (self.max+self.min) / 2, math.sqrt(1/12*(self.max-self.min)*(self.max-self.min))
+
+        if self.distribution.startswith("meta"):
+            self.hparams = {}
+            def sample_meta(f):
+                indicators, passed = unpack_dict_of_tuples({hp: self.hparams[hp]() for hp in self.hparams})
+                # sampled_embeddings = list(itertools.chain.from_iterable([sampled_embeddings[k] for k in sampled_embeddings]))
+                meta_passed = f(**passed)
+                return indicators, meta_passed
+
+            args_passed = {'device': device, 'embedding_dim': embedding_dim}
+            if self.distribution == "meta_beta":
+                ## Truncated normal where std and mean are drawn randomly logarithmically scaled
+                if hasattr(self, 'b') and hasattr(self, 'k'):
+                    self.hparams = {'b': lambda: (None, self.b), 'k': lambda: (None, self.k)}
+                else:
+                    self.hparams = {"b": DifferentiableHyperparameter(distribution="uniform", min=self.min
+                                                                                          , max=self.max, **args_passed)
+                                    , "k": DifferentiableHyperparameter(distribution="uniform", min=self.min
+                                                                                           , max=self.max, **args_passed)}
+                def make_beta(b, k):
+                    return lambda b=b, k=k: self.scale * beta_sampler_f(b, k)()
+                self.sampler = lambda make_beta=make_beta : sample_meta(make_beta)
+            elif self.distribution == "meta_trunc_norm_log_scaled":
+                # these choices are copied down below, don't change these without changing `replace_differentiable_distributions`
+                self.min_std = self.min_std if hasattr(self, 'min_std') else 0.001
+                self.max_std = self.max_std if hasattr(self, 'max_std') else self.max_mean
+                ## Truncated normal where std and mean are drawn randomly logarithmically scaled
+                if not hasattr(self, 'log_mean'):
+                    self.hparams = {"log_mean": DifferentiableHyperparameter(distribution="uniform", min=math.log(self.min_mean)
+                                                                                          , max=math.log(self.max_mean), **args_passed)
+                                    , "log_std": DifferentiableHyperparameter(distribution="uniform", min=math.log(self.min_std)
+                                                                                           , max=math.log(self.max_std), **args_passed)}
+                else:
+                    self.hparams = {'log_mean': lambda: (None, self.log_mean), 'log_std': lambda: (None, self.log_std)}
+                def make_trunc_norm(log_mean, log_std):
+                    return ((lambda : self.lower_bound + round(trunc_norm_sampler_f(math.exp(log_mean), math.exp(log_std))())) if self.round
+                            else (lambda: self.lower_bound + trunc_norm_sampler_f(math.exp(log_mean), math.exp(log_std))()))
+
+                self.sampler = lambda make_trunc_norm=make_trunc_norm: sample_meta(make_trunc_norm)
+            elif self.distribution == "meta_trunc_norm":
+                self.min_std = self.min_std if hasattr(self, 'min_std') else 0
+                self.max_std = self.max_std if hasattr(self, 'max_std') else self.max_mean
+                self.hparams = {"mean": DifferentiableHyperparameter(distribution="uniform", min=self.min_mean
+                                                                                      , max=self.max_mean, **args_passed)
+                                , "std": DifferentiableHyperparameter(distribution="uniform", min=self.min_std
+                                                                                       , max=self.max_std, **args_passed)}
+                def make_trunc_norm(mean, std):
+                    return ((lambda: self.lower_bound + round(
+                        trunc_norm_sampler_f(math.exp(mean), math.exp(std))())) if self.round
+                            else (
+                        lambda make_trunc_norm=make_trunc_norm: self.lower_bound + trunc_norm_sampler_f(math.exp(mean), math.exp(std))()))
+                self.sampler = lambda : sample_meta(make_trunc_norm)
+            elif self.distribution == "meta_choice":
+                if hasattr(self, 'choice_1_weight'):
+                    self.hparams = {f'choice_{i}_weight': lambda: (None, getattr(self, f'choice_{i}_weight')) for i in range(1, len(self.choice_values))}
+                else:
+                    self.hparams = {f"choice_{i}_weight": DifferentiableHyperparameter(distribution="uniform", min=-5.0
+                                                                                          , max=6.0, **args_passed) for i in range(1, len(self.choice_values))}
+                def make_choice(**choices):
+                    weights = torch.softmax(torch.tensor([1.0] + [choices[i] for i in choices], dtype=torch.float), 0)  # create a tensor of weights
+                    sample = torch.multinomial(weights, 1, replacement=True).numpy()[0]
+                    return self.choice_values[sample]
+
+                self.sampler = lambda make_choice=make_choice: sample_meta(make_choice)
+            elif self.distribution == "meta_choice_mixed":
+                if hasattr(self, 'choice_1_weight'):
+                    self.hparams = {f'choice_{i}_weight': lambda: (None, getattr(self, f'choice_{i}_weight')) for i in range(1, len(self.choice_values))}
+                else:
+                    self.hparams = {f"choice_{i}_weight": DifferentiableHyperparameter(distribution="uniform", min=-5.0
+                                                                                          , max=6.0, **args_passed) for i in range(1, len(self.choice_values))}
+                def make_choice(**choices):
+                    weights = torch.softmax(torch.tensor([1.0] + [choices[i] for i in choices], dtype=torch.float), 0)  # create a tensor of weights
+                    def sample():
+                        s = torch.multinomial(weights, 1, replacement=True).numpy()[0]
+                        return self.choice_values[s]()
+                    return lambda: sample
+
+                self.sampler = lambda make_choice=make_choice: sample_meta(make_choice)
+        else:
+            def return_two(x, min, max, mean, std):
+                # Returns (a hyperparameter value, and an indicator value passed to the model)
+                if mean is not None:
+                    ind = (x-mean)/std#(2 * (x-min) / (max-min) - 1)
+                else:
+                    ind = None
+                return ind, x # normalize indicator to [-1, 1]
+            # def sample_standard(sampler_f, embedding):
+            #     s = torch.tensor([sampler_f()], device = self.device)
+            #     return s, embedding(s)
+            self.sampler_f, self.sampler_min, self.sampler_max, self.sampler_mean, self.sampler_std = get_sampler()
+            self.sampler = lambda : return_two(self.sampler_f(), min=self.sampler_min, max=self.sampler_max
+                                               , mean=self.sampler_mean, std=self.sampler_std)
+            # self.embedding_layer = nn.Linear(1, self.embedding_dim, device=self.device)
+            # self.embed = lambda x : self.embedding_layer(
+            #     (x - self.sampler_min) / (self.sampler_max - self.sampler_min))
+            #self.sampler = lambda : sample_standard(self.sampler_f, self.embedding)
+
+
+    def forward(self):
+        s, s_passed = self.sampler()
+        return s, s_passed
+
+
+
+class DifferentiableHyperparameterList(nn.Module):
+    def __init__(self, hyperparameters, embedding_dim, device):
+        super().__init__()
+
+        self.device = device
+        hyperparameters = {k: v for (k, v) in hyperparameters.items() if v}
+        self.hyperparameters = nn.ModuleDict({hp: DifferentiableHyperparameter(embedding_dim = embedding_dim
+                                                                               , name = hp
+                                                                               , device = device, **hyperparameters[hp]) for hp in hyperparameters})
+    def get_hyperparameter_info(self):
+        sampled_hyperparameters_f, sampled_hyperparameters_keys = [], []
+        def append_hp(hp_key, hp_val):
+            sampled_hyperparameters_keys.append(hp_key)
+            # Function remaps hyperparameters from [-1, 1] range to true value
+            s_min, s_max, s_mean, s_std = hp_val.sampler_min, hp_val.sampler_max, hp_val.sampler_mean, hp_val.sampler_std
+            sampled_hyperparameters_f.append((lambda x: (x-s_mean)/s_std, lambda y : (y * s_std)+s_mean))
+            #sampled_hyperparameters_f.append(((lambda x: ((x - s_min) / (s_max - s_min) * (2) - 1)
+            #                                  , (lambda y: ((y + 1) * (1 / 2) * (s_max - s_min) + s_min))))
+        for hp in self.hyperparameters:
+            hp_val = self.hyperparameters[hp]
+            if hasattr(hp_val, 'hparams'):
+                for hp_ in hp_val.hparams:
+                    append_hp(f'{hp}_{hp_}', hp_val.hparams[hp_])
+            else:
+                append_hp(hp, hp_val)
+
+
+        return sampled_hyperparameters_keys, sampled_hyperparameters_f
+
+    def sample_parameter_object(self):
+        sampled_hyperparameters, s_passed = {}, {}
+        for hp in self.hyperparameters:
+            sampled_hyperparameters_, s_passed_ = self.hyperparameters[hp]()
+            s_passed[hp] = s_passed_
+            if isinstance(sampled_hyperparameters_, dict):
+                sampled_hyperparameters_ = {hp + '_' + str(key): val for key, val in sampled_hyperparameters_.items()}
+                sampled_hyperparameters.update(sampled_hyperparameters_)
+            else:
+                sampled_hyperparameters[hp] = sampled_hyperparameters_
+
+        # s_passed contains the values passed to the get_batch function
+        # sampled_hyperparameters contains the indicator of the sampled value, i.e. only number that describe the sampled object
+        return s_passed, sampled_hyperparameters#self.pack_parameter_object(sampled_embeddings)
+
+class DifferentiablePrior(torch.nn.Module):
+    def __init__(self, get_batch, hyperparameters, differentiable_hyperparameters, args):
+        super(DifferentiablePrior, self).__init__()
+
+        self.h = hyperparameters
+        self.args = args
+        self.get_batch = get_batch
+        self.differentiable_hyperparameters = DifferentiableHyperparameterList(differentiable_hyperparameters
+                                                                               , embedding_dim=self.h['emsize']
+                                                                               , device=self.args['device'])
+
+    def forward(self):
+        # Sample hyperparameters
+        sampled_hyperparameters_passed, sampled_hyperparameters_indicators = self.differentiable_hyperparameters.sample_parameter_object()
+
+        hyperparameters = {**self.h, **sampled_hyperparameters_passed}
+        x, y, y_ = self.get_batch(hyperparameters=hyperparameters, **self.args)
+
+        return x, y, y_, sampled_hyperparameters_indicators
+
+
+# TODO: Make this a class that keeps objects
+@torch.no_grad()
+def get_batch(batch_size, seq_len, num_features, get_batch
+              , device=default_device, differentiable_hyperparameters={}
+              , hyperparameters=None, batch_size_per_gp_sample=None, **kwargs):
+    batch_size_per_gp_sample = batch_size_per_gp_sample or (min(64, batch_size))
+    num_models = batch_size // batch_size_per_gp_sample
+    assert num_models * batch_size_per_gp_sample == batch_size, f'Batch size ({batch_size}) not divisible by batch_size_per_gp_sample ({batch_size_per_gp_sample})'
+
+    args = {'device': device, 'seq_len': seq_len, 'num_features': num_features, 'batch_size': batch_size_per_gp_sample}
+
+    models = [DifferentiablePrior(get_batch, hyperparameters, differentiable_hyperparameters, args) for _ in range(num_models)]
+    sample = sum([[model()] for model in models], [])
+
+    x, y, y_, hyperparameter_dict = zip(*sample)
+
+    if 'verbose' in hyperparameters and hyperparameters['verbose']:
+        print('Hparams', hyperparameter_dict[0].keys())
+
+    hyperparameter_matrix = []
+    for batch in hyperparameter_dict:
+        hyperparameter_matrix.append([batch[hp] for hp in batch])
+
+    transposed_hyperparameter_matrix = list(zip(*hyperparameter_matrix))
+    assert all([all([hp is None for hp in hp_]) or all([hp is not None for hp in hp_]) for hp_ in transposed_hyperparameter_matrix]), 'it should always be the case that when a hyper-parameter is None, once it is always None'
+    # we remove columns that are only None (i.e. not sampled)
+    hyperparameter_matrix = [[hp for hp in hp_ if hp is not None] for hp_ in hyperparameter_matrix]
+    if len(hyperparameter_matrix[0]) > 0:
+        packed_hyperparameters = torch.tensor(hyperparameter_matrix)
+        packed_hyperparameters = torch.repeat_interleave(packed_hyperparameters, repeats=batch_size_per_gp_sample, dim=0).detach()
+    else:
+        packed_hyperparameters = None
+
+    x, y, y_, packed_hyperparameters = (torch.cat(x, 1).detach()
+                                        , torch.cat(y, 1).detach()
+                                        , torch.cat(y_, 1).detach()
+                                        , packed_hyperparameters)#list(itertools.chain.from_iterable(itertools.repeat(x, batch_size_per_gp_sample) for x in packed_hyperparameters)))#torch.repeat_interleave(torch.stack(packed_hyperparameters, 0).detach(), repeats=batch_size_per_gp_sample, dim=0))
+
+    return x, y, y_, packed_hyperparameters
+
+DataLoader = get_batch_to_dataloader(get_batch)
+DataLoader.num_outputs = 1
+#DataLoader.validate = lambda : 0
+
+def draw_random_style(dl, device):
+    (hp_embedding, data, targets_), targets = next(iter(dl))
+    return hp_embedding.to(device)[0:1, :]
+
+def merge_style_with_info(diff_hparams_keys, diff_hparams_f, style, transform=True):
+    params = dict(zip(diff_hparams_keys, zip(diff_hparams_f, style.detach().cpu().numpy().tolist()[0])))
+    def t(v):
+        if transform:
+            return v[0][1](v[1])
+        else:
+            return v[1]
+    return {k : t(v) for k, v in params.items()}
+
+
+import ConfigSpace.hyperparameters as CSH
+
+def replace_differentiable_distributions(config):
+    diff_config = config['differentiable_hyperparameters']
+    for name, diff_hp_dict in diff_config.items():
+        distribution = diff_hp_dict['distribution']
+        if distribution == 'uniform':
+            diff_hp_dict['sample'] = CSH.UniformFloatHyperparameter(name, diff_hp_dict['min'], diff_hp_dict['max'])
+        elif distribution == 'meta_beta':
+            diff_hp_dict['k'] = CSH.UniformFloatHyperparameter(name+'_k', diff_hp_dict['min'], diff_hp_dict['max'])
+            diff_hp_dict['b'] = CSH.UniformFloatHyperparameter(name+'_b', diff_hp_dict['min'], diff_hp_dict['max'])
+        elif distribution == 'meta_choice':
+            for i in range(1, len(diff_hp_dict['choice_values'])):
+                diff_hp_dict[f'choice_{i}_weight'] = CSH.UniformFloatHyperparameter(name+f'choice_{i}_weight', -5.0, 6.0)
+        elif distribution == 'meta_choice_mixed':
+            for i in range(1, len(diff_hp_dict['choice_values'])):
+                diff_hp_dict[f'choice_{i}_weight'] = CSH.UniformFloatHyperparameter(name+f'choice_{i}_weight', -5.0, 6.0)
+        elif distribution == 'meta_trunc_norm_log_scaled':
+            diff_hp_dict['log_mean'] = CSH.UniformFloatHyperparameter(name+'_log_mean', math.log(diff_hp_dict['min_mean']), math.log(diff_hp_dict['max_mean']))
+            min_std = diff_hp_dict['min_std'] if 'min_std' in diff_hp_dict else 0.001
+            max_std = diff_hp_dict['max_std'] if 'max_std' in diff_hp_dict else diff_hp_dict['max_mean']
+            diff_hp_dict['log_std'] = CSH.UniformFloatHyperparameter(name+'_log_std', math.log(min_std), math.log(max_std))
+        else:
+            raise ValueError(f'Unknown distribution {distribution}')
+

TabPFN/priors/fast_gp.py

@@ -0,0 +1,144 @@
+import time
+
+import torch
+from torch import nn
+import gpytorch
+
+from .utils import get_batch_to_dataloader
+from utils import default_device
+
+
+# We will use the simplest form of GP model, exact inference
+class ExactGPModel(gpytorch.models.ExactGP):
+    def __init__(self, train_x, train_y, likelihood):
+        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
+        self.mean_module = gpytorch.means.ConstantMean()
+        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
+
+    def forward(self, x):
+        mean_x = self.mean_module(x)
+        covar_x = self.covar_module(x)
+        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
+
+
+def get_model(x, y, hyperparameters):
+    likelihood = gpytorch.likelihoods.GaussianLikelihood(noise_constraint=gpytorch.constraints.GreaterThan(1.e-9))
+    model = ExactGPModel(x, y, likelihood)
+    model.likelihood.noise = torch.ones_like(model.likelihood.noise) * hyperparameters["noise"]
+    model.covar_module.outputscale = torch.ones_like(model.covar_module.outputscale) * hyperparameters["outputscale"]
+    model.covar_module.base_kernel.lengthscale = torch.ones_like(model.covar_module.base_kernel.lengthscale) * \
+                                                 hyperparameters["lengthscale"]
+    return model, likelihood
+
+
+@torch.no_grad()
+def get_batch(batch_size, seq_len, num_features, device=default_device, hyperparameters=None,
+              equidistant_x=False, fix_x=None, **kwargs):
+    if isinstance(hyperparameters, (tuple, list)):
+        hyperparameters = {"noise": hyperparameters[0]
+            , "outputscale": hyperparameters[1]
+            , "lengthscale": hyperparameters[2]
+            , "is_binary_classification": hyperparameters[3]
+            # , "num_features_used": hyperparameters[4]
+            , "normalize_by_used_features": hyperparameters[5]
+            , "order_y": hyperparameters[6]
+            , "sampling": hyperparameters[7]
+                           }
+    elif hyperparameters is None:
+        hyperparameters = {"noise": .1, "outputscale": .1, "lengthscale": .1}
+
+    if 'verbose' in hyperparameters and hyperparameters['verbose']:
+        print({"noise": hyperparameters['noise'], "outputscale": hyperparameters['outputscale']
+                  , "lengthscale": hyperparameters['lengthscale'], 'batch_size': batch_size, 'sampling': hyperparameters['sampling']})
+
+    # hyperparameters = {k: hyperparameters[k]() if callable(hyperparameters[k]) else hyperparameters[k] for k in
+    #      hyperparameters.keys()}
+    assert not (equidistant_x and (fix_x is not None))
+
+    with gpytorch.settings.fast_computations(*hyperparameters.get('fast_computations', (True, True, True))):
+        if equidistant_x:
+            assert num_features == 1
+            x = torch.linspace(0, 1., seq_len).unsqueeze(0).repeat(batch_size, 1).unsqueeze(-1)
+        elif fix_x is not None:
+            assert fix_x.shape == (seq_len, num_features)
+            x = fix_x.unsqueeze(0).repeat(batch_size, 1, 1).to(device)
+        else:
+            if hyperparameters.get('sampling','uniform') == 'uniform':
+                x = torch.rand(batch_size, seq_len, num_features, device=device)
+            else:
+                x = torch.randn(batch_size, seq_len, num_features, device=device)
+        model, likelihood = get_model(x, torch.Tensor(), hyperparameters)
+        model.to(device)
+        # trained_model = ExactGPModel(train_x, train_y, likelihood).cuda()
+        # trained_model.eval()
+        is_fitted = False
+        while not is_fitted:
+            try:
+                with gpytorch.settings.prior_mode(True):
+                    model, likelihood = get_model(x, torch.Tensor(), hyperparameters)
+                    model.to(device)
+
+                    d = model(x)
+                    d = likelihood(d)
+                    sample = d.sample().transpose(0, 1)
+                    is_fitted = True
+            except RuntimeError: # This can happen when torch.linalg.eigh fails. Restart with new init resolves this.
+                print('GP Fitting unsuccessful, retrying.. ')
+                print(x)
+                print(hyperparameters)
+
+    if bool(torch.any(torch.isnan(x)).detach().cpu().numpy()):
+        print({"noise": hyperparameters['noise'], "outputscale": hyperparameters['outputscale']
+                  , "lengthscale": hyperparameters['lengthscale'], 'batch_size': batch_size})
+
+    # TODO: Multi output
+    return x.transpose(0, 1), sample, sample  # x.shape = (T,B,H)
+
+DataLoader = get_batch_to_dataloader(get_batch)
+DataLoader.num_outputs = 1
+
+def get_model_on_device(x,y,hyperparameters,device):
+    model, likelihood = get_model(x, y, hyperparameters)
+    model.to(device)
+    return model, likelihood
+
+
+@torch.no_grad()
+def evaluate(x, y, y_non_noisy, use_mse=False, hyperparameters={}, get_model_on_device=get_model_on_device, device=default_device, step_size=1, start_pos=0):
+    start_time = time.time()
+    losses_after_t = [.0] if start_pos == 0 else []
+    all_losses_after_t = []
+
+    with gpytorch.settings.fast_computations(*hyperparameters.get('fast_computations',(True,True,True))), gpytorch.settings.fast_pred_var(False):
+        for t in range(max(start_pos, 1), len(x), step_size):
+            loss_sum = 0.
+            model, likelihood = get_model_on_device(x[:t].transpose(0, 1), y[:t].transpose(0, 1), hyperparameters, device)
+
+
+            model.eval()
+            # print([t.shape for t in model.train_inputs])
+            # print(x[:t].transpose(0,1).shape, x[t].unsqueeze(1).shape, y[:t].transpose(0,1).shape)
+            f = model(x[t].unsqueeze(1))
+            l = likelihood(f)
+            means = l.mean.squeeze()
+            varis = l.covariance_matrix.squeeze()
+            # print(l.variance.squeeze(), l.mean.squeeze(), y[t])
+
+            assert len(means.shape) == len(varis.shape) == 1
+            assert len(means) == len(varis) == x.shape[1]
+
+            if use_mse:
+                c = nn.MSELoss(reduction='none')
+                ls = c(means, y[t])
+            else:
+                ls = -l.log_prob(y[t].unsqueeze(1))
+
+            losses_after_t.append(ls.mean())
+            all_losses_after_t.append(ls.flatten())
+        return torch.stack(all_losses_after_t).to('cpu'), torch.tensor(losses_after_t).to('cpu'), time.time() - start_time
+
+if __name__ == '__main__':
+    hps = (.1,.1,.1)
+    for redo_idx in range(1):
+        print(
+            evaluate(*get_batch(1000, 10, hyperparameters=hps, num_features=10), use_mse=False, hyperparameters=hps))

TabPFN/priors/flexible_categorical.py

@@ -0,0 +1,240 @@
+import time
+import random
+
+import torch
+from torch import nn
+
+from .utils import get_batch_to_dataloader
+from utils import normalize_data, nan_handling_missing_for_unknown_reason_value, nan_handling_missing_for_no_reason_value, nan_handling_missing_for_a_reason_value, to_ranking_low_mem, remove_outliers
+from .utils import normalize_by_used_features_f, randomize_classes, CategoricalActivation
+from .utils import uniform_int_sampler_f
+
+time_it = False
+
+class BalancedBinarize(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return (x > torch.median(x)).float()
+
+def class_sampler_f(min_, max_):
+    def s():
+        if random.random() > 0.5:
+            return uniform_int_sampler_f(min_, max_)()
+        return 2
+    return s
+
+class MulticlassRank(nn.Module):
+    def __init__(self, num_classes, ordered_p=0.5):
+        super().__init__()
+        self.num_classes = class_sampler_f(2, num_classes)()
+        self.ordered_p = ordered_p
+
+    def forward(self, x):
+        # x has shape (T,B,H)
+
+        # CAUTION: This samples the same idx in sequence for each class boundary in a batch
+        class_boundaries = torch.randint(0, x.shape[0], (self.num_classes - 1,))
+        class_boundaries = x[class_boundaries].unsqueeze(1)
+
+        d = (x > class_boundaries).sum(axis=0)
+
+        randomized_classes = torch.rand((d.shape[1], )) > self.ordered_p
+        d[:, randomized_classes] = randomize_classes(d[:, randomized_classes], self.num_classes)
+        reverse_classes = torch.rand((d.shape[1],)) > 0.5
+        d[:, reverse_classes] = self.num_classes - 1 - d[:, reverse_classes]
+        return d
+
+class MulticlassValue(nn.Module):
+    def __init__(self, num_classes, ordered_p=0.5):
+        super().__init__()
+        self.num_classes = class_sampler_f(2, num_classes)()
+        self.classes = nn.Parameter(torch.randn(num_classes-1), requires_grad=False)
+        self.ordered_p = ordered_p
+
+    def forward(self, x):
+        # x has shape (T,B,H)
+        d = (x > (self.classes.unsqueeze(-1).unsqueeze(-1))).sum(axis=0)
+
+        randomized_classes = torch.rand((d.shape[1],)) > self.ordered_p
+        d[:, randomized_classes] = randomize_classes(d[:, randomized_classes], self.num_classes)
+        reverse_classes = torch.rand((d.shape[1],)) > 0.5
+        d[:, reverse_classes] = self.num_classes - 1 - d[:, reverse_classes]
+        return d
+
+class MulticlassMultiNode(nn.Module):
+    def __init__(self, num_classes, ordered_p=0.5):
+        super().__init__()
+        self.num_classes = class_sampler_f(2, num_classes)()
+        self.classes = nn.Parameter(torch.randn(num_classes-1), requires_grad=False)
+        self.alt_multi_class = MulticlassValue(num_classes, ordered_p)
+
+    def forward(self, x):
+        # x has shape T, B, H
+        if len(x.shape) == 2:
+            return self.alt_multi_class(x)
+        T = 3
+        x[torch.isnan(x)] = 0.00001
+        d = torch.multinomial(torch.pow(0.00001+torch.sigmoid(x[:, :, 0:self.num_classes]).reshape(-1, self.num_classes), T), 1, replacement=True).reshape(x.shape[0], x.shape[1]).float()
+        return d
+
+
+class FlexibleCategorical(torch.nn.Module):
+    def __init__(self, get_batch, hyperparameters, args):
+        super(FlexibleCategorical, self).__init__()
+
+        self.h = {k: hyperparameters[k]() if callable(hyperparameters[k]) else hyperparameters[k] for k in
+                                hyperparameters.keys()}
+        self.args = args
+        self.args_passed = {**self.args}
+        self.args_passed.update({'num_features': self.h['num_features_used']})
+        self.get_batch = get_batch
+
+        if self.h['num_classes'] > 1 and not self.h['balanced']:
+            if self.h['multiclass_type'] == 'rank':
+                self.class_assigner = MulticlassRank(self.h['num_classes']
+                                             , ordered_p=self.h['output_multiclass_ordered_p']
+                                             )
+            elif self.h['multiclass_type'] == 'value':
+                self.class_assigner = MulticlassValue(self.h['num_classes']
+                                                     , ordered_p=self.h['output_multiclass_ordered_p']
+                                                     )
+            elif self.h['multiclass_type'] == 'multi_node':
+                self.class_assigner = MulticlassMultiNode(self.h['num_classes'])
+            else:
+                raise ValueError("Unknow Multiclass type")
+        elif self.h['num_classes'] == 2 and self.h['balanced']:
+            self.class_assigner = BalancedBinarize()
+        elif self.h['num_classes'] > 2 and self.h['balanced']:
+            raise NotImplementedError("Balanced multiclass training is not possible")
+        else:
+            self.class_assigner = lambda x:x # Regression
+
+    def drop_for_reason(self, x, v):
+        nan_prob_sampler = CategoricalActivation(ordered_p=0.0
+                                                 , categorical_p=1.0
+                                                 , keep_activation_size=False,
+                                                 num_classes_sampler=lambda: 20)
+        d = nan_prob_sampler(x)
+        # TODO: Make a different ordering for each activation
+        x[d < torch.rand((1,), device=x.device) * 20 * self.h['nan_prob_no_reason'] * random.random()] = v
+        return x
+
+    def drop_for_no_reason(self, x, v):
+        x[torch.rand(x.shape, device=self.args['device']) < self.h['nan_prob_no_reason']] = v
+        return x
+
+    def forward(self, batch_size):
+        start = time.time()
+        x, y, y_ = self.get_batch(hyperparameters=self.h, **self.args_passed)
+        if time_it:
+            print('Flex Forward Block 1', round(time.time() - start, 3))
+
+        start = time.time()
+
+        if self.h['nan_prob_no_reason']+self.h['nan_prob_a_reason']+self.h['nan_prob_unknown_reason'] > 0 and random.random() > 0.5: # Only one out of two datasets should have nans
+            if self.h['nan_prob_no_reason'] > 0 and random.random() > 0.5: # Missing for no reason
+                x = self.drop_for_no_reason(x, nan_handling_missing_for_no_reason_value(self.h['set_value_to_nan']))
+
+            if self.h['nan_prob_a_reason'] > 0 and random.random() > 0.5: # Missing for a reason
+                x = self.drop_for_reason(x, nan_handling_missing_for_a_reason_value(self.h['set_value_to_nan']))
+
+            if self.h['nan_prob_unknown_reason'] > 0: # Missing for unknown reason  and random.random() > 0.5
+                if random.random() < self.h['nan_prob_unknown_reason_reason_prior']:
+                    x = self.drop_for_no_reason(x, nan_handling_missing_for_unknown_reason_value(self.h['set_value_to_nan']))
+                else:
+                    x = self.drop_for_reason(x, nan_handling_missing_for_unknown_reason_value(self.h['set_value_to_nan']))
+
+        # Categorical features
+        if 'categorical_feature_p' in self.h and random.random() > 1 - self.h['categorical_feature_p']:
+            p = random.random()
+            for col in range(x.shape[2]):
+                m = MulticlassRank(10, ordered_p=0.3)
+                if random.random() > p:
+                    x[:, :, col] = m(x[:, :, col])
+
+        if time_it:
+            print('Flex Forward Block 2', round(time.time() - start, 3))
+            start = time.time()
+
+        if self.h['normalize_to_ranking']:
+            x = to_ranking_low_mem(x)
+        else:
+            x = remove_outliers(x)
+        x, y = normalize_data(x), normalize_data(y)
+
+        if time_it:
+            print('Flex Forward Block 3', round(time.time() - start, 3))
+            start = time.time()
+
+        # Cast to classification if enabled
+        y = self.class_assigner(y).float()
+
+        if time_it:
+            print('Flex Forward Block 4', round(time.time() - start, 3))
+            start = time.time()
+        if self.h['normalize_by_used_features']:
+            x = normalize_by_used_features_f(x, self.h['num_features_used'], self.args['num_features'], normalize_with_sqrt=self.h.get('normalize_with_sqrt',False))
+        if time_it:
+            print('Flex Forward Block 5', round(time.time() - start, 3))
+
+        start = time.time()
+        # Append empty features if enabled
+        x = torch.cat(
+            [x, torch.zeros((x.shape[0], x.shape[1], self.args['num_features'] - self.h['num_features_used']),
+                            device=self.args['device'])], -1)
+        if time_it:
+            print('Flex Forward Block 6', round(time.time() - start, 3))
+
+        return x, y, y  # x.shape = (T,B,H)
+
+import torch.cuda as cutorch
+
+@torch.no_grad()
+def get_batch(batch_size, seq_len, num_features, get_batch, device, hyperparameters=None, batch_size_per_gp_sample=None, **kwargs):
+    batch_size_per_gp_sample = batch_size_per_gp_sample or (min(32, batch_size))
+    num_models = batch_size // batch_size_per_gp_sample
+    assert num_models > 0, f'Batch size ({batch_size}) is too small for batch_size_per_gp_sample ({batch_size_per_gp_sample})'
+    assert num_models * batch_size_per_gp_sample == batch_size, f'Batch size ({batch_size}) not divisible by batch_size_per_gp_sample ({batch_size_per_gp_sample})'
+
+    # Sample one seq_len for entire batch
+    seq_len = hyperparameters['seq_len_used']() if callable(hyperparameters['seq_len_used']) else seq_len
+
+    args = {'device': device, 'seq_len': seq_len, 'num_features': num_features, 'batch_size': batch_size_per_gp_sample}
+
+    models = [FlexibleCategorical(get_batch, hyperparameters, args).to(device) for _ in range(num_models)]
+
+    start = time.time()
+    sample = sum([[model(batch_size=batch_size_per_gp_sample)] for model in models], [])
+    #print('sample', time.time() - start)
+
+    x, y, y_ = zip(*sample)
+    x, y, y_ = torch.cat(x, 1).detach(), torch.cat(y, 1).detach(), torch.cat(y_, 1).detach()
+
+    # # TODO: Reintegrate this code (Doesn't work on batch dim), could be applied to each batch sample individually
+    # if hyperparameters['is_binary_classification'] and hyperparameters['order_y']:
+    #     x, y = order_by_y(x, y)
+
+    return x, y, y_
+
+# num_features_used = num_features_used_sampler()
+# prior_outputscale = prior_outputscale_sampler()
+# prior_lengthscale = prior_lengthscale_sampler()
+#
+# x, sample = normalize_data(x), normalize_data(sample)
+#
+# if is_binary_classification:
+#     sample = (sample > torch.median(sample, dim=0)[0]).float()
+#
+# if normalize_by_used_features:
+#     x = normalize_by_used_features_f(x, num_features_used, num_features)
+#
+# # # if is_binary_classification and order_y:
+# # #     x, sample = order_by_y(x, sample)
+# #
+# # Append empty features if enabled
+# x = torch.cat([x, torch.zeros((x.shape[0], x.shape[1], num_features - num_features_used), device=device)], -1)
+
+DataLoader = get_batch_to_dataloader(get_batch)
+DataLoader.num_outputs = 1

TabPFN/priors/mlp.py

@@ -0,0 +1,173 @@
+import random
+import math
+
+import torch
+from torch import nn
+import numpy as np
+
+from utils import default_device
+from .utils import get_batch_to_dataloader
+
+class GaussianNoise(nn.Module):
+    def __init__(self, std, device):
+        super().__init__()
+        self.std = std
+        self.device=device
+
+    def forward(self, x):
+        return x + torch.normal(torch.zeros_like(x), self.std)
+
+
+def causes_sampler_f(num_causes):
+    means = np.random.normal(0, 1, (num_causes))
+    std = np.abs(np.random.normal(0, 1, (num_causes)) * means)
+    return means, std
+
+def get_batch(batch_size, seq_len, num_features, hyperparameters, device=default_device, num_outputs=1, sampling='normal', **kwargs):
+    if ('mix_activations' in hyperparameters) and hyperparameters['mix_activations']:
+        s = hyperparameters['prior_mlp_activations']()
+        hyperparameters['prior_mlp_activations'] = lambda : s
+
+    class MLP(torch.nn.Module):
+        def __init__(self, hyperparameters):
+            super(MLP, self).__init__()
+
+            with torch.no_grad():
+
+                for key in hyperparameters:
+                    setattr(self, key, hyperparameters[key])
+
+                assert (self.num_layers >= 2)
+
+                if 'verbose' in hyperparameters and self.verbose:
+                    print({k : hyperparameters[k] for k in ['is_causal', 'num_causes', 'prior_mlp_hidden_dim'
+                        , 'num_layers', 'noise_std', 'y_is_effect', 'pre_sample_weights', 'prior_mlp_dropout_prob'
+                        , 'pre_sample_causes']})
+
+                if self.is_causal:
+                    self.prior_mlp_hidden_dim = max(self.prior_mlp_hidden_dim, num_outputs + 2 * num_features)
+                else:
+                    self.num_causes = num_features
+
+                # This means that the mean and standard deviation of each cause is determined in advance
+                if self.pre_sample_causes:
+                    self.causes_mean, self.causes_std = causes_sampler_f(self.num_causes)
+                    self.causes_mean = torch.tensor(self.causes_mean, device=device).unsqueeze(0).unsqueeze(0).tile(
+                        (seq_len, 1, 1))
+                    self.causes_std = torch.tensor(self.causes_std, device=device).unsqueeze(0).unsqueeze(0).tile(
+                        (seq_len, 1, 1))
+
+                def generate_module(layer_idx, out_dim):
+                    # Determine std of each noise term in initialization, so that is shared in runs
+                    # torch.abs(torch.normal(torch.zeros((out_dim)), self.noise_std)) - Change std for each dimension?
+                    noise = (GaussianNoise(torch.abs(torch.normal(torch.zeros(size=(1, out_dim), device=device), float(self.noise_std))), device=device)
+                         if self.pre_sample_weights else GaussianNoise(float(self.noise_std), device=device))
+                    return [
+                        nn.Sequential(*[self.prior_mlp_activations()
+                            , nn.Linear(self.prior_mlp_hidden_dim, out_dim)
+                            , noise])
+                    ]
+
+                self.layers = [nn.Linear(self.num_causes, self.prior_mlp_hidden_dim, device=device)]
+                self.layers += [module for layer_idx in range(self.num_layers-1) for module in generate_module(layer_idx, self.prior_mlp_hidden_dim)]
+                if not self.is_causal:
+                    self.layers += generate_module(-1, num_outputs)
+                self.layers = nn.Sequential(*self.layers)
+
+                # Initialize Model parameters
+                for i, (n, p) in enumerate(self.layers.named_parameters()):
+                    if self.block_wise_dropout:
+                        if len(p.shape) == 2: # Only apply to weight matrices and not bias
+                            nn.init.zeros_(p)
+                            # TODO: N blocks should be a setting
+                            n_blocks = random.randint(1, math.ceil(math.sqrt(min(p.shape[0], p.shape[1]))))
+                            w, h = p.shape[0] // n_blocks, p.shape[1] // n_blocks
+                            keep_prob = (n_blocks*w*h) / p.numel()
+                            for block in range(0, n_blocks):
+                                nn.init.normal_(p[w * block: w * (block+1), h * block: h * (block+1)], std=self.init_std / keep_prob**(1/2))
+                    else:
+                        if len(p.shape) == 2: # Only apply to weight matrices and not bias
+                            dropout_prob = self.prior_mlp_dropout_prob if i > 0 else 0.0  # Don't apply dropout in first layer
+                            dropout_prob = min(dropout_prob, 0.99)
+                            nn.init.normal_(p, std=self.init_std / (1. - dropout_prob)**(1/2))
+                            p *= torch.bernoulli(torch.zeros_like(p) + 1. - dropout_prob)
+
+        def forward(self):
+            def sample_normal():
+                if self.pre_sample_causes:
+                    causes = torch.normal(self.causes_mean, self.causes_std.abs()).float()
+                else:
+                    causes = torch.normal(0., 1., (seq_len, 1, self.num_causes), device=device).float()
+                return causes
+
+            if self.sampling == 'normal':
+                causes = sample_normal()
+            elif self.sampling == 'mixed':
+                zipf_p, multi_p, normal_p = random.random() * 0.66, random.random() * 0.66, random.random() * 0.66
+                def sample_cause(n):
+                    if random.random() > normal_p:
+                        if self.pre_sample_causes:
+                            return torch.normal(self.causes_mean[:, :, n], self.causes_std[:, :, n].abs()).float()
+                        else:
+                            return torch.normal(0., 1., (seq_len, 1), device=device).float()
+                    elif random.random() > multi_p:
+                        x = torch.multinomial(torch.rand((random.randint(2, 10))), seq_len, replacement=True).to(device).unsqueeze(-1).float()
+                        x = (x - torch.mean(x)) / torch.std(x)
+                        return x
+                    else:
+                        x = torch.minimum(torch.tensor(np.random.zipf(2.0 + random.random() * 2, size=(seq_len)),
+                                            device=device).unsqueeze(-1).float(), torch.tensor(10.0, device=device))
+                        return x - torch.mean(x)
+                causes = torch.cat([sample_cause(n).unsqueeze(-1) for n in range(self.num_causes)], -1)
+            elif self.sampling == 'uniform':
+                causes = torch.rand((seq_len, 1, self.num_causes), device=device)
+            else:
+                raise ValueError(f'Sampling is set to invalid setting: {sampling}.')
+
+            outputs = [causes]
+            for layer in self.layers:
+                outputs.append(layer(outputs[-1]))
+            outputs = outputs[2:]
+
+            if self.is_causal:
+                ## Sample nodes from graph if model is causal
+                outputs_flat = torch.cat(outputs, -1)
+
+                if self.in_clique:
+                    random_perm = random.randint(0, outputs_flat.shape[-1] - num_outputs - num_features) + torch.randperm(num_outputs + num_features, device=device)
+                else:
+                    random_perm = torch.randperm(outputs_flat.shape[-1]-1, device=device)
+
+                random_idx_y = list(range(-num_outputs, -0)) if self.y_is_effect else random_perm[0:num_outputs]
+                random_idx = random_perm[num_outputs:num_outputs + num_features]
+
+                if self.sort_features:
+                    random_idx, _ = torch.sort(random_idx)
+                y = outputs_flat[:, :, random_idx_y]
+
+                x = outputs_flat[:, :, random_idx]
+            else:
+                y = outputs[-1][:, :, :]
+                x = causes
+
+            if bool(torch.any(torch.isnan(x)).detach().cpu().numpy()) or bool(torch.any(torch.isnan(y)).detach().cpu().numpy()):
+                x[:] = 0.0
+                y[:] = 1.0
+
+            return x, y
+
+    model = MLP(hyperparameters).to(device)
+
+    sample = sum([[model()] for _ in range(0, batch_size)], [])
+
+    x, y = zip(*sample)
+    y = torch.cat(y, 1).detach().squeeze(2)
+    x = torch.cat(x, 1).detach()
+    x = x[..., torch.randperm(x.shape[-1])]
+
+    return x, y, y
+
+
+DataLoader = get_batch_to_dataloader(get_batch)
+DataLoader.num_outputs = 1
+

TabPFN/priors/prior.py

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

TabPFN/priors/prior_bag.py

@@ -0,0 +1,32 @@
+import torch
+
+from .utils import get_batch_to_dataloader
+from utils import default_device
+
+def get_batch(batch_size, seq_len, num_features, device=default_device
+              , hyperparameters=None, batch_size_per_gp_sample=None, **kwargs):
+    batch_size_per_gp_sample = batch_size_per_gp_sample or (min(64, batch_size))
+    num_models = batch_size // batch_size_per_gp_sample
+    assert num_models * batch_size_per_gp_sample == batch_size, f'Batch size ({batch_size}) not divisible by batch_size_per_gp_sample ({batch_size_per_gp_sample})'
+
+    args = {'device': device, 'seq_len': seq_len, 'num_features': num_features, 'batch_size': batch_size_per_gp_sample}
+
+    prior_bag_priors_get_batch = hyperparameters['prior_bag_get_batch']
+    prior_bag_priors_p = [1.0] + [hyperparameters[f'prior_bag_exp_weights_{i}'] for i in range(1, len(prior_bag_priors_get_batch))]
+
+    weights = torch.tensor(prior_bag_priors_p, dtype=torch.float)  # create a tensor of weights
+    batch_assignments = torch.multinomial(torch.softmax(weights, 0), num_models, replacement=True).numpy()
+
+    if 'verbose' in hyperparameters and hyperparameters['verbose']:
+        print('PRIOR_BAG:', weights, batch_assignments)
+
+    sample = sum([[prior_bag_priors_get_batch[int(prior_idx)](hyperparameters=hyperparameters, **args)] for prior_idx in batch_assignments], [])
+
+    x, y, y_ = zip(*sample)
+    x, y, y_ = (torch.cat(x, 1).detach()
+                                        , torch.cat(y, 1).detach()
+                                        , torch.cat(y_, 1).detach())
+    return x, y, y_
+
+DataLoader = get_batch_to_dataloader(get_batch)
+DataLoader.num_outputs = 1

TabPFN/priors/utils.py

@@ -0,0 +1,163 @@
+import random
+
+import torch
+
+from utils import set_locals_in_self
+from .prior import PriorDataLoader
+from torch import nn
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.gridspec as gridspec
+import scipy.stats as stats
+import math
+
+def get_batch_to_dataloader(get_batch_method_):
+    class DL(PriorDataLoader):
+        get_batch_method = get_batch_method_
+
+        # Caution, you might need to set self.num_features manually if it is not part of the args.
+        def __init__(self, num_steps, fuse_x_y=False, **get_batch_kwargs):
+            set_locals_in_self(locals())
+            # The stuff outside the or is set as class attribute before instantiation.
+            self.num_features = get_batch_kwargs.get('num_features') or self.num_features
+            self.num_outputs = get_batch_kwargs.get('num_outputs') or self.num_outputs
+            print('DataLoader.__dict__', self.__dict__)
+
+        @staticmethod
+        def gbm(*args, fuse_x_y=True, **kwargs):
+            dynamic_seq_len = callable(kwargs['seq_len'])
+            kwargs['seq_len'] = kwargs['seq_len']() if dynamic_seq_len else kwargs['seq_len']
+            # Scales the batch size dynamically with the power of 'dynamic_batch_size'.
+            # A transformer with quadratic memory usage in the seq len would need a power of 2 to keep memory constant.
+            if dynamic_seq_len and 'dynamic_batch_size' in kwargs and kwargs['dynamic_batch_size'] > 0:
+                kwargs['batch_size'] = kwargs['batch_size'] * math.floor(math.pow(kwargs['seq_len_maximum'], kwargs['dynamic_batch_size']) / math.pow(kwargs['seq_len'], kwargs['dynamic_batch_size']))
+            batch = get_batch_method_(*args, **kwargs)
+            x, y, target_y, style = batch if len(batch) == 4 else (batch[0], batch[1], batch[2], None)
+            if fuse_x_y:
+                return torch.cat([x, torch.cat([torch.zeros_like(y[:1]), y[:-1]], 0).unsqueeze(-1).float()],
+                                 -1), target_y
+            else:
+                return (style, x, y), target_y
+
+        def __len__(self):
+            return self.num_steps
+
+        def __iter__(self):
+            return iter(self.gbm(**self.get_batch_kwargs, fuse_x_y=self.fuse_x_y) for _ in range(self.num_steps))
+
+
+    return DL
+
+import seaborn as sns
+def plot_features(data, targets, fig=None):
+    if torch.is_tensor(data):
+        data = data.detach().cpu().numpy()
+        targets = targets.detach().cpu().numpy()
+    #data = np.concatenate([data, data[:, -1:]], -1)
+    #df = pd.DataFrame(data, columns=list(range(0, data.shape[1])))
+    #g = sns.pairplot(df, hue=data.shape[1]-1, palette="Set2", diag_kind="kde", height=2.5)
+    #plt.legend([], [], frameon=False)
+    #g._legend.remove()
+    #g = sns.PairGrid(df, hue=data.shape[1]-1)
+    #g.map_diag(sns.histplot)
+    #g.map_offdiag(sns.scatterplot)
+    #g._legend.remove()
+
+    fig2 = fig if fig else plt.figure(figsize=(8, 8))
+    spec2 = gridspec.GridSpec(ncols=data.shape[1], nrows=data.shape[1], figure=fig2)
+    for d in range(0, data.shape[1]):
+        for d2 in range(0, data.shape[1]):
+            sub_ax = fig2.add_subplot(spec2[d, d2])
+            if d == d2:
+                sns.kdeplot(data[:, d],hue=targets[:],ax=sub_ax,legend=False, palette="deep")
+                sub_ax.set(ylabel=None)
+            else:
+                sns.scatterplot(x=data[:, d], y=data[:, d2],
+                           hue=targets[:],legend=False, palette="deep")
+                #plt.scatter(data[:, d], data[:, d2],
+                #               c=targets[:])
+            sub_ax.get_xaxis().set_ticks([])
+            sub_ax.get_yaxis().set_ticks([])
+    plt.subplots_adjust(wspace=0.05, hspace=0.05)
+    fig2.show()
+
+
+def plot_prior(prior):
+    s = np.array([prior() for _ in range(0, 1000)])
+    count, bins, ignored = plt.hist(s, 50, density=True)
+    print(s.min())
+    plt.show()
+
+trunc_norm_sampler_f = lambda mu, sigma : lambda: stats.truncnorm((0 - mu) / sigma, (1000000 - mu) / sigma, loc=mu, scale=sigma).rvs(1)[0]
+beta_sampler_f = lambda a, b : lambda : np.random.beta(a, b)
+gamma_sampler_f = lambda a, b : lambda : np.random.gamma(a, b)
+uniform_sampler_f = lambda a, b : lambda : np.random.uniform(a, b)
+uniform_int_sampler_f = lambda a, b : lambda : round(np.random.uniform(a, b))
+def zipf_sampler_f(a, b, c):
+    x = np.arange(b, c)
+    weights = x ** (-a)
+    weights /= weights.sum()
+    return lambda : stats.rv_discrete(name='bounded_zipf', values=(x, weights)).rvs(1)
+scaled_beta_sampler_f = lambda a, b, scale, minimum : lambda : minimum + round(beta_sampler_f(a, b)() * (scale - minimum))
+
+
+def normalize_by_used_features_f(x, num_features_used, num_features, normalize_with_sqrt=False):
+    if normalize_with_sqrt:
+        return x / (num_features_used / num_features)**(1 / 2)
+    return x / (num_features_used / num_features)
+
+
+def order_by_y(x, y):
+    order = torch.argsort(y if random.randint(0, 1) else -y, dim=0)[:, 0, 0]
+    order = order.reshape(2, -1).transpose(0, 1).reshape(-1)#.reshape(seq_len)
+    x = x[order]  # .reshape(2, -1).transpose(0, 1).reshape(-1).flip([0]).reshape(seq_len, 1, -1)
+    y = y[order]  # .reshape(2, -1).transpose(0, 1).reshape(-1).reshape(seq_len, 1, -1)
+
+    return x, y
+
+def randomize_classes(x, num_classes):
+    classes = torch.arange(0, num_classes, device=x.device)
+    random_classes = torch.randperm(num_classes, device=x.device).type(x.type())
+    x = ((x.unsqueeze(-1) == classes) * random_classes).sum(-1)
+    return x
+
+
+class CategoricalActivation(nn.Module):
+    def __init__(self, categorical_p=0.1, ordered_p=0.7
+                 , keep_activation_size=False
+                 , num_classes_sampler=zipf_sampler_f(0.8, 1, 10)):
+        self.categorical_p = categorical_p
+        self.ordered_p = ordered_p
+        self.keep_activation_size = keep_activation_size
+        self.num_classes_sampler = num_classes_sampler
+
+        super().__init__()
+
+    def forward(self, x):
+        # x shape: T, B, H
+
+        x = nn.Softsign()(x)
+
+        num_classes = self.num_classes_sampler()
+        hid_strength = torch.abs(x).mean(0).unsqueeze(0) if self.keep_activation_size else None
+
+        categorical_classes = torch.rand((x.shape[1], x.shape[2])) < self.categorical_p
+        class_boundaries = torch.zeros((num_classes - 1, x.shape[1], x.shape[2]), device=x.device, dtype=x.dtype)
+        # Sample a different index for each hidden dimension, but shared for all batches
+        for b in range(x.shape[1]):
+            for h in range(x.shape[2]):
+                ind = torch.randint(0, x.shape[0], (num_classes - 1,))
+                class_boundaries[:, b, h] = x[ind, b, h]
+
+        for b in range(x.shape[1]):
+            x_rel = x[:, b, categorical_classes[b]]
+            boundaries_rel = class_boundaries[:, b, categorical_classes[b]].unsqueeze(1)
+            x[:, b, categorical_classes[b]] = (x_rel > boundaries_rel).sum(dim=0).float() - num_classes / 2
+
+        ordered_classes = torch.rand((x.shape[1],x.shape[2])) < self.ordered_p
+        ordered_classes = torch.logical_and(ordered_classes, categorical_classes)
+        x[:, ordered_classes] = randomize_classes(x[:, ordered_classes], num_classes)
+
+        x = x * hid_strength if self.keep_activation_size else x
+
+        return x

TabPFN/requirements.txt

@@ -0,0 +1,15 @@
+# Please use python V 3.7 to be compatible with all packages
+gpytorch==1.5.0
+torch==1.9.0
+scikit-learn==0.24.2
+pyyaml==5.4.1
+seaborn==0.11.2
+xgboost==1.4.0
+tqdm==4.62.1
+numpy==1.21.2
+openml==0.12.2
+catboost==0.26.1
+auto-sklearn==0.14.5
+hyperopt==0.2.5
+configspace==0.4.21
+# autogluon==0.4.0

TabPFN/scripts/baseline_prediction_interface.py

@@ -0,0 +1,38 @@
+import tqdm
+import numpy as np
+
+def baseline_predict(metric_function, eval_xs, eval_ys, categorical_feats, metric_used=None, eval_pos=2, max_time=300, **kwargs):
+    """
+    Baseline prediction interface.
+    :param metric_function:
+    :param eval_xs:
+    :param eval_ys:
+    :param categorical_feats:
+    :param metric_used:
+    :param eval_pos:
+    :param max_time: Scheduled maximum time
+    :param kwargs:
+    :return: list [np.array(metrics), np.array(outputs), best_configs] or [None, None, None] if failed
+    """
+
+    metrics = []
+    outputs = []
+    best_configs = []
+    eval_splits = list(zip(eval_xs.transpose(0, 1), eval_ys.transpose(0, 1)))
+    for eval_x, eval_y in tqdm.tqdm(eval_splits, desc='Calculating splits'+str(metric_function)+' '+str(eval_pos)):
+        try:
+            metric, output, best_config = metric_function(eval_x[:eval_pos],
+                                                          eval_y[:eval_pos],
+                                                          eval_x[eval_pos:],
+                                                          eval_y[eval_pos:],
+                                                          categorical_feats,
+                                                          metric_used=metric_used
+                                                          , max_time=max_time)
+            metrics += [metric]
+            outputs += [output]
+            best_configs += [best_config]
+            return np.array(metrics), np.array(outputs), best_configs
+        except Exception as e:
+            print(f'There was an exception in {metric_function}')
+            print(e)
+            return None, None, None

TabPFN/scripts/differentiable_pfn_evaluation.py

@@ -0,0 +1,391 @@
+import os
+import torch
+import numpy as np
+import time
+import pickle
+from scripts import  tabular_metrics
+from scripts.tabular_metrics import calculate_score_per_method
+from scripts.tabular_evaluation import evaluate
+from priors.differentiable_prior import draw_random_style
+from tqdm import tqdm
+from pathlib import Path
+import random
+from model_builder import load_model
+from scripts.transformer_prediction_interface import get_params_from_config
+
+"""
+===============================
+PUBLIC FUNCTIONS FOR EVALUATION
+===============================
+"""
+
+
+def eval_model_range(i_range, *args, **kwargs):
+    for i in i_range:
+        eval_model(i, *args, **kwargs)
+
+
+def load_model_workflow(i, e, add_name, base_path, device='cpu', eval_addition=''):
+    """
+    Workflow for loading a model and setting appropriate parameters for diffable hparam tuning.
+
+    :param i:
+    :param e:
+    :param eval_positions_valid:
+    :param add_name:
+    :param base_path:
+    :param device:
+    :param eval_addition:
+    :return:
+    """
+    def check_file(e):
+        model_file = f'models_diff/prior_diff_real_checkpoint{add_name}_n_{i}_epoch_{e}.cpkt'
+        model_path = os.path.join(base_path, model_file)
+        # print('Evaluate ', model_path)
+        results_file = os.path.join(base_path,
+                                    f'models_diff/prior_diff_real_results{add_name}_n_{i}_epoch_{e}_{eval_addition}.pkl')
+        if not Path(model_path).is_file():  # or Path(results_file).is_file():
+            return None, None, None
+        return model_file, model_path, results_file
+
+    model_file = None
+    if e == -1:
+        for e_ in range(100, -1, -1):
+            model_file_, model_path_, results_file_ = check_file(e_)
+            if model_file_ is not None:
+                e = e_
+                model_file, model_path, results_file = model_file_, model_path_, results_file_
+                break
+    else:
+        model_file, model_path, results_file = check_file(e)
+
+    if model_file is None:
+        print('No checkpoint found')
+        return None
+
+    print(f'Loading {model_file}')
+
+    model, c = load_model(base_path, model_file, device, eval_positions=[], verbose=False)
+
+    return model, c, results_file
+
+
+def eval_model(i, e, valid_datasets, test_datasets, train_datasets, eval_positions_valid, eval_positions_test,
+               bptt_valid,
+               bptt_test, add_name, base_path, device='cpu', eval_addition='', **extra_tuning_args):
+    """
+    Differentiable model evaliation workflow. Evaluates and saves results to disk.
+
+    :param i:
+    :param e:
+    :param valid_datasets:
+    :param test_datasets:
+    :param train_datasets:
+    :param eval_positions_valid:
+    :param eval_positions_test:
+    :param bptt_valid:
+    :param bptt_test:
+    :param add_name:
+    :param base_path:
+    :param device:
+    :param eval_addition:
+    :param extra_tuning_args:
+    :return:
+    """
+    model, c, results_file = load_model_workflow(i, e, add_name, base_path, device, eval_addition)
+    params = {'bptt': bptt_valid
+        , 'bptt_final': bptt_test
+        , 'eval_positions': eval_positions_valid
+        , 'eval_positions_test': eval_positions_test
+        , 'valid_datasets': valid_datasets
+        , 'test_datasets': test_datasets
+        , 'train_datasets': train_datasets
+        , 'verbose': True
+        , 'device': device
+              }
+
+    params.update(get_params_from_config(c))
+
+    start = time.time()
+    metrics, metrics_valid, style, temperature, optimization_route = evaluate_differentiable_model(model, **params,
+                                                                                                   **extra_tuning_args)
+    print('Evaluation time: ', time.time() - start)
+
+    print(results_file)
+    r = [c.copy(), metrics, metrics_valid, style.to('cpu'), temperature.to('cpu'), optimization_route]
+    with open(results_file, 'wb') as output:
+        del r[0]['num_features_used']
+        del r[0]['categorical_features_sampler']
+        pickle.dump(r, output)
+
+    _, _, _, style, temperature, _ = r
+
+    return r, model
+
+"""
+===============================
+INTERNAL HELPER FUNCTIONS
+===============================
+"""
+
+def evaluate_differentiable_model(model
+                                  , valid_datasets
+                                  , test_datasets
+                                  , train_datasets
+                                  , N_draws=100
+                                  , N_grad_steps=10
+                                  , eval_positions=None
+                                  , eval_positions_test=None
+                                  , bptt=100
+                                  , bptt_final=200
+                                  , style=None
+                                  , n_parallel_configurations=1
+                                  , device='cpu'
+                                  , selection_metric='auc'
+                                  , final_splits=[1, 2, 3, 4, 5]
+                                  , N_ensemble_configurations_list=[1, 5, 10, 20, 50, 100]
+                                  , **kwargs):
+    """
+    Evaluation function for diffable model evaluation. Returns a list of results.
+
+    :param model:
+    :param valid_datasets:
+    :param test_datasets:
+    :param train_datasets:
+    :param N_draws:
+    :param N_grad_steps:
+    :param eval_positions:
+    :param eval_positions_test:
+    :param bptt:
+    :param bptt_final:
+    :param style:
+    :param n_parallel_configurations:
+    :param device:
+    :param selection_metric:
+    :param final_splits:
+    :param N_ensemble_configurations_list:
+    :param kwargs:
+    :return:
+    """
+    torch.manual_seed(0)
+    np.random.seed(0)
+    random.seed(0)
+
+    diffable_metric = tabular_metrics.cross_entropy
+    evaluation_metric = tabular_metrics.auc_metric
+    if selection_metric in ('auc', 'roc'):
+        selection_metric_min_max = 'max'
+        selection_metric = tabular_metrics.auc_metric
+        evaluation_metric = selection_metric
+    elif selection_metric in ('ce', 'selection_metric'):
+        selection_metric_min_max = 'min'
+        selection_metric = tabular_metrics.cross_entropy
+        evaluation_metric = selection_metric
+
+    print('Diffable metric', diffable_metric, ' Selection metric', selection_metric, ' Evaluation metric',
+          evaluation_metric)
+    print('N PARALLEL CONFIGURATIONS', n_parallel_configurations)
+    print('eval_positions', eval_positions)
+
+    def evaluate_valid(style, softmax_temperature, results, results_tracked):
+        result_valid = eval_step(valid_datasets, style, softmax_temperature=softmax_temperature,
+                                 return_tensor=False, inference_mode=True, selection_metric=selection_metric,
+                                 evaluation_metric=evaluation_metric, eval_positions=eval_positions, bptt=bptt, model=model[2])
+        result_valid = [float(result_valid[f'mean_select_at_{pos}']) for pos in eval_positions]
+        results += [result_valid]
+        results_tracked += [np.nanmean(result_valid)]
+
+    model[2].to(device)
+    model[2].eval()
+
+    results_on_valid, results_on_valid_tracked = [], []
+    best_style, best_softmax_temperature = style, torch.cat(
+        [torch.tensor([0.0]).to(device) for n in range(0, n_parallel_configurations)], 0)
+    optimization_routes = []
+
+    best_style = torch.cat([draw_random_style(model[3], device).detach() for n in range(0, n_parallel_configurations)],
+                      0)
+    best_softmax_temperature = torch.cat([torch.tensor([0.0]).to(device) for n in range(0, n_parallel_configurations)],
+                                    0)
+
+
+    for _ in tqdm(range(0, N_draws), desc='Iterate over Optimization initializations'): # Evaluates N hparam draws
+        style = torch.cat([draw_random_style(model[3], device).detach() for n in range(0, n_parallel_configurations)],
+                          0)
+        softmax_temperature = torch.cat([torch.tensor([0.0]).to(device) for n in range(0, n_parallel_configurations)],
+                                        0)
+
+        evaluate_valid(style, softmax_temperature, results_on_valid, results_on_valid_tracked)
+
+        print(f'Draw --> Valid Selection metric: {results_on_valid[-1]}')
+
+        if N_grad_steps > 0:
+            gradient_optimize_result = gradient_optimize_style(model, style, N_grad_steps
+                                                               , softmax_temperature=softmax_temperature
+                                                               , model=model[2]
+                                                               , train_datasets=train_datasets
+                                                               , valid_datasets=valid_datasets
+                                                               , selection_metric_min_max=selection_metric_min_max
+                                                               , **kwargs)
+            optimization_routes += [gradient_optimize_result['optimization_route']]
+
+            evaluate_valid(gradient_optimize_result['best_style']
+                                          , gradient_optimize_result['best_temperature']
+                                          , results_on_valid, results_on_valid_tracked)
+
+            print(f'After diff --> Valid Selection metric: {results_on_valid[-1]}')
+
+        if selection_metric_min_max == 'min':
+            is_best = (results_on_valid_tracked[-1] <= min(results_on_valid_tracked))
+        else:
+            is_best = (results_on_valid_tracked[-1] >= max(results_on_valid_tracked))
+
+        if is_best or best_style is None:
+            best_style = gradient_optimize_result['best_style'].clone()
+            best_softmax_temperature = gradient_optimize_result['best_temperature'].clone()
+    torch.cuda.empty_cache()
+
+    def final_evaluation():
+        print('Running eval dataset with final params (no gradients)..')
+        print(best_style, best_softmax_temperature)
+        result_test = []
+        for N_ensemble_configurations in N_ensemble_configurations_list:
+            print(f'Running with {N_ensemble_configurations} ensemble_configurations')
+            kwargs['N_ensemble_configurations'] = N_ensemble_configurations
+            splits = []
+            for split in final_splits:
+                splits += [eval_step(test_datasets, best_style, softmax_temperature=best_softmax_temperature
+                                     , return_tensor=False, eval_positions=eval_positions_test,
+                                     bptt=bptt_final, inference_mode=True, split_number=split, model=model[2]
+                                     , selection_metric=selection_metric, evaluation_metric=evaluation_metric)]
+            result_test += [splits]
+
+        print('Running valid dataset with final params (no gradients)..')
+        result_valid = eval_step(valid_datasets, best_style, softmax_temperature=best_softmax_temperature
+                                 , return_tensor=False, eval_positions=eval_positions_test,
+                                 bptt=bptt_final, inference_mode=True, model=model[2]
+                                 , selection_metric=selection_metric, evaluation_metric=evaluation_metric)
+
+        return result_test, result_valid
+
+    result_test, result_valid = final_evaluation()
+
+    return result_test, result_valid, best_style, best_softmax_temperature, optimization_routes
+
+
+def eval_step(ds, used_style, selection_metric, evaluation_metric, eval_positions, return_tensor=True, **kwargs):
+    def step():
+        return evaluate(datasets=ds,
+                        method='transformer'
+                        , overwrite=True
+                        , style=used_style
+                        , eval_positions=eval_positions
+                        , metric_used=selection_metric
+                        , save=False
+                        , path_interfix=None
+                        , base_path=None
+                        , verbose=True
+                        , **kwargs)
+
+    if return_tensor:
+        r = step()
+    else:
+        with torch.no_grad():
+            r = step()
+
+    calculate_score_per_method(selection_metric, 'select', r, ds, eval_positions, aggregator='mean')
+    calculate_score_per_method(evaluation_metric, 'eval', r, ds, eval_positions, aggregator='mean')
+
+    return r
+
+
+def gradient_optimize_style(model, init_style, steps, softmax_temperature, train_datasets, valid_datasets, learning_rate=0.03, optimize_all=False,
+                            limit_style=True, N_datasets_sampled=90, optimize_softmax_temperature=True, selection_metric_min_max='max', **kwargs):
+    """
+    Uses gradient based methods to optimize 'style' on the 'train_datasets' and uses stopping with 'valid_datasets'.
+
+    :param model:
+    :param init_style:
+    :param steps:
+    :param learning_rate:
+    :param softmax_temperature:
+    :param train_datasets:
+    :param valid_datasets:
+    :param optimize_all:
+    :param limit_style:
+    :param N_datasets_sampled:
+    :param optimize_softmax_temperature:
+    :param selection_metric_min_max:
+    :param kwargs:
+    :return:
+    """
+    grad_style = torch.nn.Parameter(init_style.detach(), requires_grad=True)
+
+    best_style, best_temperature, best_selection_metric, best_diffable_metric = grad_style.detach(), softmax_temperature.detach(), None, None
+    softmax_temperature = torch.nn.Parameter(softmax_temperature.detach(), requires_grad=optimize_softmax_temperature)
+    variables_to_optimize = model[2].parameters() if optimize_all else [grad_style, softmax_temperature]
+    optimizer = torch.optim.Adam(variables_to_optimize, lr=learning_rate)
+
+    optimization_route_selection, optimization_route_diffable = [], []
+    optimization_route_selection_valid, optimization_route_diffable_valid = [], []
+
+    def eval_opt(ds, return_tensor=True, inference_mode=False):
+        result = eval_step(ds, grad_style, softmax_temperature=softmax_temperature, return_tensor=return_tensor
+                           , inference_mode=inference_mode, model=model[2], **kwargs)
+
+        diffable_metric = result['mean_metric']
+        selection_metric = result['mean_select']
+
+        return diffable_metric, selection_metric
+
+    def eval_all_datasets(datasets, propagate=True):
+        selection_metrics_this_step, diffable_metrics_this_step = [], []
+        for ds in datasets:
+            diffable_metric_train, selection_metric_train = eval_opt([ds], inference_mode=(not propagate))
+            if not torch.isnan(diffable_metric_train).any():
+                if propagate and diffable_metric_train.requires_grad == True:
+                    diffable_metric_train.backward()
+                selection_metrics_this_step += [selection_metric_train]
+                diffable_metrics_this_step += [float(diffable_metric_train.detach().cpu().numpy())]
+        diffable_metric_train = np.nanmean(diffable_metrics_this_step)
+        selection_metric_train = np.nanmean(selection_metrics_this_step)
+
+        return diffable_metric_train, selection_metric_train
+
+    for t in tqdm(range(steps), desc='Iterate over Optimization steps'):
+        optimizer.zero_grad()
+
+        # Select subset of datasets
+        random.seed(t)
+        train_datasets_ = random.sample(train_datasets, N_datasets_sampled)
+
+        # Get score on train
+        diffable_metric_train, selection_metric_train = eval_all_datasets(train_datasets_, propagate=True)
+        optimization_route_selection += [float(selection_metric_train)]
+        optimization_route_diffable += [float(diffable_metric_train)]
+
+        # Get score on valid
+        diffable_metric_valid, selection_metric_valid = eval_all_datasets(valid_datasets, propagate=False)
+        optimization_route_selection_valid += [float(selection_metric_valid)]
+        optimization_route_diffable_valid += [float(diffable_metric_valid)]
+
+        is_best = (selection_metric_min_max == 'min' and best_selection_metric > selection_metric_valid)
+        is_best = is_best or (selection_metric_min_max == 'max' and best_selection_metric < selection_metric_valid)
+        if (best_selection_metric is None) or (not np.isnan(selection_metric_valid) and is_best):
+            print('New best', best_selection_metric, selection_metric_valid)
+            best_style = grad_style.detach().clone()
+            best_temperature = softmax_temperature.detach().clone()
+            best_selection_metric, best_diffable_metric = selection_metric_valid, diffable_metric_valid
+
+        optimizer.step()
+
+        if limit_style:
+            grad_style = grad_style.detach().clamp(-1.74, 1.74)
+
+        print(f'Valid: Diffable metric={diffable_metric_valid} Selection metric={selection_metric_valid};' +
+            f'Train: Diffable metric={diffable_metric_train} Selection metric={selection_metric_train}')
+
+    print(f'Return best:{best_style} {best_selection_metric}')
+    return {'best_style': best_style, 'best_temperature': best_temperature
+            , 'optimization_route': {'select': optimization_route_selection, 'loss': optimization_route_diffable,
+               'test_select': optimization_route_selection_valid, 'test_loss': optimization_route_diffable_valid}}

TabPFN/scripts/model_configs.py

@@ -0,0 +1,210 @@
+from copy import deepcopy
+from priors.utils import uniform_int_sampler_f
+from priors.differentiable_prior import DifferentiableHyperparameter
+from ConfigSpace import hyperparameters as CSH
+import torch
+from priors.differentiable_prior import replace_differentiable_distributions
+
+import ConfigSpace as CS
+
+def get_general_config(max_features, bptt, eval_positions=None):
+    """"
+    Returns the general PFN training hyperparameters.
+    """
+    config_general = {
+        "lr": CSH.UniformFloatHyperparameter('lr', lower=0.00002, upper=0.0002, log=True),
+        "dropout": CSH.CategoricalHyperparameter('dropout', [0.0]),
+        "emsize": CSH.CategoricalHyperparameter('emsize', [2 ** i for i in range(8, 9)]), ## upper bound is -1
+        "batch_size": CSH.CategoricalHyperparameter('batch_size', [2 ** i for i in range(8, 9)]),
+        "nlayers": CSH.CategoricalHyperparameter('nlayers', [12]),
+        "num_features": max_features,
+        "nhead": CSH.CategoricalHyperparameter('nhead', [4]),
+        "nhid_factor": 2,
+        "bptt": bptt,
+        "eval_positions": None,
+        "seq_len_used": bptt,
+        "sampling": 'normal',#hp.choice('sampling', ['mixed', 'normal']), # uniform
+        "epochs": 80,
+        "num_steps": 100,
+        "verbose": False,
+        "pre_sample_causes": True, # This is MLP
+        "mix_activations": False,#hp.choice('mix_activations', [True, False]),
+    }
+
+    return config_general
+
+def get_flexible_categorical_config(max_features):
+    """"
+    Returns the configuration parameters for the tabular multiclass wrapper.
+    """
+    config_flexible_categorical = {
+        "nan_prob_unknown_reason_reason_prior": CSH.CategoricalHyperparameter('nan_prob_unknown_reason_reason_prior', [1.0]),
+        "categorical_feature_p": CSH.CategoricalHyperparameter('categorical_feature_p', [0.0]),
+        "nan_prob_no_reason": CSH.CategoricalHyperparameter('nan_prob_no_reason', [0.0, 0.1, 0.2]),
+        "nan_prob_unknown_reason": CSH.CategoricalHyperparameter('nan_prob_unknown_reason', [0.0]),
+        "nan_prob_a_reason": CSH.CategoricalHyperparameter('nan_prob_a_reason', [0.0]),
+        # "num_classes": lambda : random.randint(2, 10), "balanced": False,
+        "max_num_classes": 2,
+        "num_classes": 2,
+        "noise_type": CSH.CategoricalHyperparameter('noise_type', ["Gaussian"]), # NN
+        "balanced": True,
+        "normalize_to_ranking": CSH.CategoricalHyperparameter('normalize_to_ranking', [False]),
+        "set_value_to_nan": CSH.CategoricalHyperparameter('set_value_to_nan', [0.5, 0.2, 0.0]),
+        "normalize_by_used_features": True,
+        "num_features_used":
+            {'uniform_int_sampler_f(3,max_features)': uniform_int_sampler_f(1, max_features)}
+        # hp.choice('conv_activation', [{'distribution': 'uniform', 'min': 2.0, 'max': 8.0}, None]),
+    }
+    return config_flexible_categorical
+
+def get_diff_flex():
+    """"
+    Returns the configuration parameters for a differentiable wrapper around the tabular multiclass wrapper.
+    """
+    diff_flex = {
+        # "ordinal_pct": {'distribution': 'uniform', 'min': 0.0, 'max': 0.5},
+        # "num_categorical_features_sampler_a": hp.choice('num_categorical_features_sampler_a',
+        #                                                 [{'distribution': 'uniform', 'min': 0.3, 'max': 0.9}, None]),
+        # "num_categorical_features_sampler_b": {'distribution': 'uniform', 'min': 0.3, 'max': 0.9},
+        "output_multiclass_ordered_p": {'distribution': 'uniform', 'min': 0.0, 'max': 0.5}, #CSH.CategoricalHyperparameter('output_multiclass_ordered_p', [0.0, 0.1, 0.2]),
+        "multiclass_type": {'distribution': 'meta_choice', 'choice_values': ['value', 'rank']},
+    }
+
+    return diff_flex
+
+def get_diff_gp():
+    """"
+    Returns the configuration parameters for a differentiable wrapper around GP.
+    """
+    diff_gp = {
+        'outputscale': {'distribution': 'meta_trunc_norm_log_scaled', 'max_mean': 10., 'min_mean': 0.00001, 'round': False,
+                        'lower_bound': 0},
+        'lengthscale': {'distribution': 'meta_trunc_norm_log_scaled', 'max_mean': 10., 'min_mean': 0.00001, 'round': False,
+                        'lower_bound': 0},
+        'noise': {'distribution': 'meta_choice', 'choice_values': [0.00001, 0.0001, 0.01]}
+    }
+
+    return diff_gp
+
+def get_diff_causal():
+    """"
+    Returns the configuration parameters for a differentiable wrapper around MLP / Causal mixture.
+    """
+    diff_causal = {
+        "num_layers": {'distribution': 'meta_trunc_norm_log_scaled', 'max_mean': 6, 'min_mean': 1, 'round': True,
+                       'lower_bound': 2},
+        # Better beta?
+        "prior_mlp_hidden_dim": {'distribution': 'meta_trunc_norm_log_scaled', 'max_mean': 130, 'min_mean': 5,
+                                 'round': True, 'lower_bound': 4},
+
+        "prior_mlp_dropout_prob": {'distribution': 'meta_beta', 'scale': 0.9, 'min': 0.1, 'max': 5.0},
+    # This mustn't be too high since activations get too large otherwise
+
+        "noise_std": {'distribution': 'meta_trunc_norm_log_scaled', 'max_mean': .3, 'min_mean': 0.0001, 'round': False,
+                      'lower_bound': 0.0},
+        "init_std": {'distribution': 'meta_trunc_norm_log_scaled', 'max_mean': 10.0, 'min_mean': 0.01, 'round': False,
+                     'lower_bound': 0.0},
+        "num_causes": {'distribution': 'meta_trunc_norm_log_scaled', 'max_mean': 12, 'min_mean': 1, 'round': True,
+                       'lower_bound': 1},
+        "is_causal": {'distribution': 'meta_choice', 'choice_values': [True, False]},
+        "pre_sample_weights": {'distribution': 'meta_choice', 'choice_values': [True, False]},
+        "y_is_effect": {'distribution': 'meta_choice', 'choice_values': [True, False]},
+        "prior_mlp_activations": {'distribution': 'meta_choice_mixed', 'choice_values': [
+            torch.nn.Tanh
+            , torch.nn.ReLU
+            , torch.nn.Identity
+            , lambda : torch.nn.LeakyReLU(negative_slope=0.1)
+            , torch.nn.ELU
+        ]},
+        "block_wise_dropout": {'distribution': 'meta_choice', 'choice_values': [True, False]},
+        "sort_features": {'distribution': 'meta_choice', 'choice_values': [True, False]},
+        "in_clique": {'distribution': 'meta_choice', 'choice_values': [True, False]},
+    }
+
+    return diff_causal
+
+def get_diff_prior_bag():
+    """"
+    Returns the configuration parameters for a GP and MLP / Causal mixture.
+    """
+    diff_prior_bag = {
+        'prior_bag_exp_weights_1': {'distribution': 'uniform', 'min': 100000., 'max': 100001.},
+        # MLP Weight (Biased, since MLP works better, 1.0 is weight for prior number 0)
+    }
+
+    return diff_prior_bag
+
+def get_diff_config():
+    """"
+    Returns the configuration parameters for a differentiable wrapper around GP and MLP / Causal mixture priors.
+    """
+    diff_prior_bag = get_diff_prior_bag()
+    diff_causal = get_diff_causal()
+    diff_gp = get_diff_gp()
+    diff_flex = get_diff_flex()
+
+    config_diff = {'differentiable_hyperparameters': {**diff_prior_bag, **diff_causal, **diff_gp, **diff_flex}}
+
+    return config_diff
+
+
+def sample_differentiable(config):
+    """"
+    Returns sampled hyperparameters from a differentiable wrapper, that is it makes a non-differentiable out of
+    differentiable.
+    """
+    # config is a dict of dicts, dicts that have a 'distribution' key are treated as distributions to be sampled
+    result = deepcopy(config)
+    del result['differentiable_hyperparameters']
+
+    for k, v in config['differentiable_hyperparameters'].items():
+        s_indicator, s_hp = DifferentiableHyperparameter(**v, embedding_dim=None,
+                                                         device=None)()  # both of these are actually not used to the best of my knowledge
+        result[k] = s_hp
+
+    return result
+
+def list_all_hps_in_nested(config):
+    """"
+    Returns a list of hyperparameters from a neszed dict of hyperparameters.
+    """
+
+    if isinstance(config, CSH.Hyperparameter):
+        return [config]
+    elif isinstance(config, dict):
+        result = []
+        for k, v in config.items():
+            result += list_all_hps_in_nested(v)
+        return result
+    else:
+        return []
+
+def create_configspace_from_hierarchical(config):
+    cs = CS.ConfigurationSpace()
+    for hp in list_all_hps_in_nested(config):
+        cs.add_hyperparameter(hp)
+    return cs
+
+def fill_in_configsample(config, configsample):
+    # config is our dict that defines config distribution
+    # configsample is a CS.Configuration
+    hierarchical_configsample = deepcopy(config)
+    for k, v in config.items():
+        if isinstance(v, CSH.Hyperparameter):
+            hierarchical_configsample[k] = configsample[v.name]
+        elif isinstance(v, dict):
+            hierarchical_configsample[k] = fill_in_configsample(v, configsample)
+    return hierarchical_configsample
+
+
+def evaluate_hypers(config, sample_diff_hps=False):
+    """"
+    Samples a hyperparameter configuration from a sampleable configuration (can be used in HP search).
+    """
+    if sample_diff_hps:
+        # I do a deepcopy here, such that the config stays the same and can still be used with diff. hps
+        config = deepcopy(config)
+        replace_differentiable_distributions(config)
+    cs = create_configspace_from_hierarchical(config)
+    cs_sample = cs.sample_configuration()
+    return fill_in_configsample(config, cs_sample)

TabPFN/scripts/tabular_baselines.py

@@ -0,0 +1,421 @@
+from catboost import CatBoostClassifier, Pool
+
+import math
+
+from sklearn.impute import SimpleImputer
+
+import xgboost as xgb
+from sklearn import neighbors
+from sklearn.gaussian_process import GaussianProcessClassifier
+from sklearn.gaussian_process.kernels import RBF
+import numpy as np
+
+from scripts import tabular_metrics
+import pandas as pd
+
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import cross_val_score
+import time
+
+from hyperopt import fmin, tpe, hp, STATUS_OK, Trials , space_eval, rand
+from sklearn.compose import ColumnTransformer
+from sklearn.preprocessing import OneHotEncoder
+from sklearn.preprocessing import MinMaxScaler
+
+import autosklearn.classification
+
+CV = 5
+MULTITHREAD = 1 # Number of threads baselines are able to use at most
+param_grid, param_grid_hyperopt = {}, {}
+
+def get_scoring_direction(metric_used):
+    # Not needed
+    if metric_used == tabular_metrics.auc_metric:
+        return -1
+    elif metric_used == tabular_metrics.cross_entropy:
+        return 1
+    else:
+        raise Exception('No scoring string found for metric')
+
+def get_scoring_string(metric_used, multiclass=True, usage="sklearn_cv"):
+    if metric_used == tabular_metrics.auc_metric:
+        if usage == 'sklearn_cv':
+            return 'roc_auc_ovo'
+        elif usage == 'autogluon':
+            return 'log_loss' # Autogluon crashes when using 'roc_auc' with some datasets usning logloss gives better scores;
+                              # We might be able to fix this, but doesn't work out of box.
+                              # File bug report? Error happens with dataset robert and fabert
+            if multiclass:
+                return 'roc_auc_ovo_macro'
+            else:
+                return 'roc_auc'
+        elif usage == 'autosklearn':
+            if multiclass:
+                return autosklearn.metrics.log_loss # roc_auc only works for binary, use logloss instead
+            else:
+                return autosklearn.metrics.roc_auc
+        elif usage == 'catboost':
+            return 'MultiClass' # Effectively LogLoss, ROC not available
+        elif usage == 'xgb':
+            return 'logloss'
+        return 'roc_auc'
+    elif metric_used == tabular_metrics.cross_entropy:
+        if usage == 'sklearn_cv':
+            return 'neg_log_loss'
+        elif usage == 'autogluon':
+            return 'log_loss'
+        elif usage == 'autosklearn':
+            return autosklearn.metrics.log_loss
+        elif usage == 'catboost':
+            return 'MultiClass' # Effectively LogLoss
+        return 'logloss'
+    else:
+        raise Exception('No scoring string found for metric')
+
+def eval_f(params, clf_, x, y, metric_used, start_time, max_time):
+    if time.time() - start_time > max_time:
+        return np.nan
+    scores = cross_val_score(clf_(**params), x, y, cv=CV, scoring=get_scoring_string(metric_used))
+
+    return -np.nanmean(scores)
+
+def preprocess_impute(x, y, test_x, test_y, impute, one_hot, standardize, cat_features=[]):
+    import warnings
+    def warn(*args, **kwargs):
+        pass
+
+    warnings.warn = warn
+
+    x, y, test_x, test_y = x.cpu().numpy(), y.cpu().long().numpy(), test_x.cpu().numpy(), test_y.cpu().long().numpy()
+
+    if impute:
+        imp_mean = SimpleImputer(missing_values=np.nan, strategy='mean')
+        imp_mean.fit(x)
+        x, test_x = imp_mean.transform(x), imp_mean.transform(test_x)
+
+    if one_hot:
+        def make_pd_from_np(x):
+            data = pd.DataFrame(x)
+            for c in cat_features:
+                data.iloc[:, c] = data.iloc[:, c].astype('int')
+            return data
+        x, test_x = make_pd_from_np(x),  make_pd_from_np(test_x)
+        transformer = ColumnTransformer(transformers=[('cat', OneHotEncoder(handle_unknown='ignore', sparse=False), cat_features)], remainder="passthrough")
+        transformer.fit(x)
+        x, test_x = transformer.transform(x), transformer.transform(test_x)
+
+    if standardize:
+        scaler = MinMaxScaler()
+        scaler.fit(x)
+        x, test_x = scaler.transform(x), scaler.transform(test_x)
+
+    return x, y, test_x, test_y
+
+## Auto Gluon
+def autogluon_metric(x, y, test_x, test_y, cat_features, metric_used, max_time=300):
+    from autogluon.tabular import TabularPredictor # Inside function so package can be sued without installation
+    x, y, test_x, test_y = preprocess_impute(x, y, test_x, test_y
+                                             , one_hot=False
+                                             , cat_features=cat_features
+                                             , impute=False
+                                             , standardize=False)
+    train_data = pd.DataFrame(np.concatenate([x, y[:, np.newaxis]], 1))
+    test_data = pd.DataFrame(np.concatenate([test_x, test_y[:, np.newaxis]], 1))
+
+    # AutoGluon automatically infers datatypes, we don't specify the categorical labels
+    predictor = TabularPredictor(
+        label=train_data.columns[-1],
+        eval_metric=get_scoring_string(metric_used, usage='autogluon', multiclass=(len(np.unique(y)) > 2)),
+        problem_type='multiclass' if len(np.unique(y)) > 2 else 'binary'
+        ## seed=int(y[:].sum()) doesn't accept seed
+    ).fit(
+        train_data=train_data,
+        time_limit=max_time,
+        presets=['best_quality']
+        # The seed is deterministic but varies for each dataset and each split of it
+    )
+
+    pred = predictor.predict_proba(test_data, as_multiclass=True).values
+
+    metric = metric_used(test_y, pred)
+
+    return metric, pred, predictor.fit_summary()
+
+## AUTO Sklearn
+def autosklearn_metric(x, y, test_x, test_y, cat_features, metric_used, max_time=300):
+    return autosklearn2_metric(x, y, test_x, test_y, cat_features, metric_used, max_time=max_time, version=1)
+
+from autosklearn.experimental.askl2 import AutoSklearn2Classifier
+from autosklearn.classification import AutoSklearnClassifier
+def autosklearn2_metric(x, y, test_x, test_y, cat_features, metric_used, max_time=300, version=2):
+    x, y, test_x, test_y = preprocess_impute(x, y, test_x, test_y
+                                             , one_hot=False
+                                             , cat_features=cat_features
+                                             , impute=False
+                                             , standardize=False)
+
+    def make_pd_from_np(x):
+        data = pd.DataFrame(x)
+        for c in cat_features:
+            data.iloc[:, c] = data.iloc[:, c].astype('category')
+        return data
+
+    x = make_pd_from_np(x)
+    test_x = make_pd_from_np(test_x)
+
+    clf_ = AutoSklearn2Classifier if version == 2 else AutoSklearnClassifier
+    clf = clf_(time_left_for_this_task=max_time,
+                                                           memory_limit=4000,
+                                                           n_jobs=MULTITHREAD,
+               seed=int(y[:].sum()),
+        # The seed is deterministic but varies for each dataset and each split of it
+               metric=get_scoring_string(metric_used, usage='autosklearn', multiclass=len(np.unique(y)) > 2))
+
+    # fit model to data
+    clf.fit(x, y)
+
+    pred = clf.predict_proba(test_x)
+    metric = metric_used(test_y, pred)
+
+    return metric, pred, None
+
+param_grid_hyperopt['logistic'] = {
+    'penalty': hp.choice('penalty', ['l1', 'l2', 'none'])
+    , 'max_iter': hp.randint('max_iter', [50, 500])
+    , 'fit_intercept': hp.choice('fit_intercept', [True, False])
+    , 'C': hp.loguniform('C', -5, math.log(5.0))}  # 'normalize': [False],
+
+def logistic_metric(x, y, test_x, test_y, cat_features, metric_used, max_time=300):
+    x, y, test_x, test_y = preprocess_impute(x, y, test_x, test_y
+                                             , one_hot=True, impute=True, standardize=True
+                                             , cat_features=cat_features)
+
+    def clf_(**params):
+        return LogisticRegression(solver='saga', tol=1e-4, n_jobs=1, **params)
+
+    start_time = time.time()
+
+    def stop(trial):
+        return time.time() - start_time > max_time, []
+
+    best = fmin(
+        fn=lambda params: eval_f(params, clf_, x, y, metric_used, start_time, max_time),
+        space=param_grid_hyperopt['logistic'],
+        algo=rand.suggest,
+        rstate=np.random.RandomState(int(y[:].sum())),
+        early_stop_fn=stop,
+        # The seed is deterministic but varies for each dataset and each split of it
+        max_evals=10000)
+    best = space_eval(param_grid_hyperopt['logistic'], best)
+
+    clf = clf_(**best)
+    clf.fit(x, y)
+
+    pred = clf.predict_proba(test_x)
+    metric = metric_used(test_y, pred)
+
+    return metric, pred, best
+
+## KNN
+param_grid_hyperopt['knn'] = {'n_neighbors': hp.randint('n_neighbors', 1,16)
+                              }
+def knn_metric(x, y, test_x, test_y, cat_features, metric_used, max_time=300):
+    x, y, test_x, test_y = preprocess_impute(x, y, test_x, test_y,
+                                             one_hot=True, impute=True, standardize=True,
+                                             cat_features=cat_features)
+
+    def clf_(**params):
+        return neighbors.KNeighborsClassifier(n_jobs=1, **params)
+
+    start_time = time.time()
+
+    def stop(trial):
+        return time.time() - start_time > max_time, []
+
+    best = fmin(
+        fn=lambda params: eval_f(params, clf_, x, y, metric_used, start_time, max_time),
+        space=param_grid_hyperopt['knn'],
+        algo=rand.suggest,
+        rstate=np.random.RandomState(int(y[:].sum())),
+        early_stop_fn=stop,
+        # The seed is deterministic but varies for each dataset and each split of it
+        max_evals=10000)
+    best = space_eval(param_grid_hyperopt['knn'], best)
+
+    clf = clf_(**best)
+    clf.fit(x, y)
+
+    pred = clf.predict_proba(test_x)
+    metric = metric_used(test_y, pred)
+
+    return metric, pred, best
+
+## GP
+param_grid_hyperopt['gp'] = {
+    'params_y_scale': hp.loguniform('params_y_scale', math.log(0.05), math.log(5.0)),
+    'params_length_scale': hp.loguniform('params_length_scale', math.log(0.1), math.log(1.0)),
+    'n_jobs': hp.choice('njobs', [1])
+}
+def gp_metric(x, y, test_x, test_y, cat_features, metric_used, max_time=300):
+    x, y, test_x, test_y = preprocess_impute(x, y, test_x, test_y,
+                                             one_hot=True, impute=True, standardize=True,
+                                             cat_features=cat_features)
+
+    def clf_(params_y_scale,params_length_scale, **params):
+        return GaussianProcessClassifier(kernel= params_y_scale * RBF(params_length_scale), **params)
+
+    start_time = time.time()
+    def stop(trial):
+        return time.time() - start_time > max_time, []
+
+
+    best = fmin(
+        fn=lambda params: eval_f(params, clf_, x, y, metric_used, start_time, max_time),
+        space=param_grid_hyperopt['gp'],
+        algo=rand.suggest,
+        rstate=np.random.RandomState(int(y[:].sum())),
+        early_stop_fn=stop,
+        # The seed is deterministic but varies for each dataset and each split of it
+        max_evals=1000)
+    best = space_eval(param_grid_hyperopt['gp'], best)
+
+    clf = clf_(**best)
+    clf.fit(x, y)
+
+    pred = clf.predict_proba(test_x)
+    metric = metric_used(test_y, pred)
+
+    return metric, pred, best
+
+
+# Catboost
+# Hyperparameter space: https://arxiv.org/pdf/2106.03253.pdf
+
+param_grid_hyperopt['catboost'] = {
+    'learning_rate': hp.loguniform('learning_rate', math.log(math.pow(math.e, -5)), math.log(1)),
+    'random_strength': hp.randint('random_strength', 1, 20),
+    'l2_leaf_reg': hp.loguniform('l2_leaf_reg', math.log(1), math.log(10)),
+    'bagging_temperature': hp.uniform('bagging_temperature', 0., 1),
+    'leaf_estimation_iterations': hp.randint('leaf_estimation_iterations', 1, 20),
+    'iterations': hp.randint('iterations', 100, 4000), # This is smaller than in paper, 4000 leads to ram overusage
+}
+
+def catboost_metric(x, y, test_x, test_y, cat_features, metric_used, max_time=300):
+    print(x)
+
+    x, y, test_x, test_y = preprocess_impute(x, y, test_x, test_y
+                                             , one_hot=False
+                                             , cat_features=cat_features
+                                             , impute=False
+                                             , standardize=False)
+
+    # Nans in categorical features must be encoded as separate class
+    x[:, cat_features], test_x[:, cat_features] = np.nan_to_num(x[:, cat_features], -1), np.nan_to_num(
+        test_x[:, cat_features], -1)
+
+    def make_pd_from_np(x):
+        data = pd.DataFrame(x)
+        for c in cat_features:
+            data.iloc[:, c] = data.iloc[:, c].astype('int')
+        return data
+
+    x = make_pd_from_np(x)
+    test_x = make_pd_from_np(test_x)
+
+    def clf_(**params):
+        return CatBoostClassifier(
+                               loss_function=get_scoring_string(metric_used, usage='catboost'),
+                               thread_count = MULTITHREAD,
+                               used_ram_limit='4gb',
+            random_seed=int(y[:].sum()),
+                               logging_level='Silent',
+                                cat_features=cat_features,
+                                  **params)
+
+    start_time = time.time()
+    def stop(trial):
+        return time.time() - start_time > max_time, []
+
+    best = fmin(
+        fn=lambda params: eval_f(params, clf_, x, y, metric_used, start_time, max_time),
+        space=param_grid_hyperopt['catboost'],
+        algo=rand.suggest,
+        rstate=np.random.RandomState(int(y[:].sum())),
+        early_stop_fn=stop,
+        # The seed is deterministic but varies for each dataset and each split of it
+        max_evals=1000)
+    best = space_eval(param_grid_hyperopt['catboost'], best)
+
+    clf = clf_(**best)
+    clf.fit(x, y)
+
+    pred = clf.predict_proba(test_x)
+    metric = metric_used(test_y, pred)
+
+    return metric, pred, best
+
+
+# XGBoost
+# Hyperparameter space: https://arxiv.org/pdf/2106.03253.pdf
+param_grid_hyperopt['xgb'] = {
+    'learning_rate': hp.loguniform('learning_rate', -7, math.log(1)),
+    'max_depth': hp.randint('max_depth', 1, 10),
+    'subsample': hp.uniform('subsample', 0.2, 1),
+    'colsample_bytree': hp.uniform('colsample_bytree', 0.2, 1),
+    'colsample_bylevel': hp.uniform('colsample_bylevel', 0.2, 1),
+    'min_child_weight': hp.loguniform('min_child_weight', -16, 5),
+    'alpha': hp.loguniform('alpha', -16, 2),
+    'lambda': hp.loguniform('lambda', -16, 2),
+    'gamma': hp.loguniform('gamma', -16, 2),
+    'n_estimators': hp.randint('n_estimators', 100, 4000), # This is smaller than in paper
+}
+
+def xgb_metric(x, y, test_x, test_y, cat_features, metric_used, max_time=300):
+    # XGB Documentation:
+    # XGB handles categorical data appropriately without using One Hot Encoding, categorical features are experimetal
+    # XGB handles missing values appropriately without imputation
+
+    x, y, test_x, test_y = preprocess_impute(x, y, test_x, test_y
+                                             , one_hot=False
+                                             , cat_features=cat_features
+                                             , impute=False
+                                             , standardize=False)
+
+    def clf_(**params):
+        return xgb.XGBClassifier(use_label_encoder=False
+                                 , nthread=1
+                                 , **params
+                                 , eval_metric=get_scoring_string(metric_used, usage='xgb') # AUC not implemented
+        )
+
+    start_time = time.time()
+    def stop(trial):
+        return time.time() - start_time > max_time, []
+
+    best = fmin(
+        fn=lambda params: eval_f(params, clf_, x, y, metric_used, start_time, max_time),
+        space=param_grid_hyperopt['xgb'],
+        algo=rand.suggest,
+        rstate=np.random.RandomState(int(y[:].sum())),
+        early_stop_fn=stop,
+        # The seed is deterministic but varies for each dataset and each split of it
+        max_evals=1000)
+    best = space_eval(param_grid_hyperopt['xgb'], best)
+
+    clf = clf_(**best)
+    clf.fit(x, y)
+
+    pred = clf.predict_proba(test_x)
+    metric = metric_used(test_y, pred)
+
+    return metric, pred, best
+
+
+clf_dict = {'gp': gp_metric
+                , 'knn': knn_metric
+                , 'catboost': catboost_metric
+                , 'xgb': xgb_metric
+                , 'logistic': logistic_metric
+           , 'autosklearn': autosklearn_metric
+             , 'autosklearn2': autosklearn2_metric
+            , 'autogluon': autogluon_metric}

TabPFN/scripts/tabular_evaluation.py

@@ -0,0 +1,284 @@
+import time
+import os
+from pathlib import Path
+
+from tqdm import tqdm
+import random
+import numpy as np
+
+from torch import nn
+
+from utils import torch_nanmean
+from datasets import *
+from model_builder import load_model
+from scripts.tabular_baselines import get_scoring_string
+from scripts import tabular_metrics
+from scripts.transformer_prediction_interface import *
+from scripts.baseline_prediction_interface import *
+"""
+===============================
+PUBLIC FUNCTIONS FOR EVALUATION
+===============================
+"""
+
+
+def eval_model(i, e, valid_datasets, test_datasets, eval_positions, bptt, add_name, base_path, device='cpu', eval_addition='', **kwargs):
+    metrics_test, config_sample, model_path = eval_model_on_ds(i, e, test_datasets, eval_positions, bptt, add_name, base_path, device=device, eval_addition=eval_addition, **kwargs)
+    metrics_valid, _, _ = eval_model_on_ds(i, e, valid_datasets, eval_positions, bptt, add_name, base_path, device=device, eval_addition=eval_addition, **kwargs)
+    return {'mean_auc_test': metrics_test['mean_roc_at_1000'], 'mean_auc_valid': metrics_valid['mean_roc_at_1000'], 'mean_ce_test': metrics_test['mean_ce_at_1000'], 'mean_ce_valid': metrics_valid['mean_ce_at_1000'], 'config_sample': config_sample, 'model_path': model_path}
+
+def eval_model_on_ds(i, e, valid_datasets, eval_positions, bptt, add_name, base_path, device='cpu', eval_addition='', **kwargs):
+
+    # How to use: evaluate_without_fitting(i,0,valid_datasets, [1024], 100000, add_name=model_string, base_path=base_path,)
+    def check_file(e):
+        model_file = f'models_diff/prior_diff_real_checkpoint{add_name}_n_{i}_epoch_{e}.cpkt'
+        model_path = os.path.join(base_path, model_file)
+        # print('Evaluate ', model_path)
+        results_file = os.path.join(base_path,
+                                    f'models_diff/prior_diff_real_results{add_name}_n_{i}_epoch_{e}_{eval_addition}.pkl')
+        if not Path(model_path).is_file():  # or Path(results_file).is_file():
+            # print('checkpoint exists: ', Path(model_file).is_file(), ', results are written:', Path(results_file).is_file())
+            return None, None, None
+        return model_file, model_path, results_file
+
+    if e == -1: # use last checkpoint, if e == -1
+        for e_ in range(100, -1, -1):
+            model_file_, model_path_, results_file_ = check_file(e_)
+            if model_file_ is not None:
+                e = e_
+                model_file, model_path, results_file = model_file_, model_path_, results_file_
+                break
+    else:
+        model_file, model_path, results_file = check_file(e)
+
+    model, config_sample = load_model(base_path, model_file, device, None, verbose=False)
+    print(model[2].style_encoder)
+
+    params = {'max_features': config_sample['num_features']
+        , 'rescale_features': config_sample["normalize_by_used_features"]
+        , 'normalize_to_ranking': config_sample["normalize_to_ranking"]
+        , 'normalize_with_sqrt': config_sample.get("normalize_with_sqrt", False)
+              }
+    metrics_valid = evaluate(datasets=valid_datasets, model=model[2], method='transformer', device=device, overwrite=True,
+                             extend_features=True
+                             # just removed the style keyword but transformer is trained with style, just empty
+                             , save=False
+                             , metric_used=tabular_metrics.cross_entropy
+                             , return_tensor=True
+                             , verbose=False
+                             , eval_positions=eval_positions
+                             , bptt=bptt
+                             , base_path=None
+                             , inference_mode=True
+                             , **params
+                             , **kwargs)
+
+    tabular_metrics.calculate_score_per_method(tabular_metrics.auc_metric, 'roc', metrics_valid, valid_datasets, eval_positions)
+    tabular_metrics.calculate_score_per_method(tabular_metrics.cross_entropy, 'ce', metrics_valid, valid_datasets, eval_positions)
+
+    return metrics_valid, config_sample, model_path
+
+
+def evaluate(datasets, bptt, eval_positions, metric_used, model
+             , verbose=False
+             , return_tensor=False
+             , **kwargs):
+    """
+    Evaluates a list of datasets for a model function.
+
+    :param datasets: List of datasets
+    :param bptt: maximum sequence length
+    :param eval_positions: List of positions where to evaluate models
+    :param verbose: If True, is verbose.
+    :param metric_used: Which metric is optimized for.
+    :param return_tensor: Wheater to return results as a pytorch.tensor or numpy, this is only relevant for transformer.
+    :param kwargs:
+    :return:
+    """
+    overall_result = {'metric_used': get_scoring_string(metric_used)
+                      , 'bptt': bptt
+                      , 'eval_positions': eval_positions}
+
+    aggregated_metric_datasets, num_datasets = torch.tensor(0.0), 0
+
+    # For each dataset
+    for [ds_name, X, y, categorical_feats, _, _] in tqdm.tqdm(datasets, desc='Iterate over datasets') if verbose else datasets:
+        dataset_bptt = min(len(X), bptt)
+        # if verbose and dataset_bptt < bptt:
+        #    print(f'Dataset too small for given sequence length, reducing to {len(X)} ({bptt})')
+
+        aggregated_metric, num = torch.tensor(0.0), 0
+        ds_result = {}
+
+        for eval_position in (eval_positions if verbose else eval_positions):
+            eval_position_real = int(dataset_bptt * 0.5) if 2 * eval_position > dataset_bptt else eval_position
+            eval_position_bptt = int(eval_position_real * 2.0)
+
+            r = evaluate_position(X, y, model=model
+                        , num_classes=len(torch.unique(y))
+                        , categorical_feats = categorical_feats
+                        , bptt = eval_position_bptt
+                        , ds_name=ds_name
+                        , eval_position = eval_position_real
+                        , metric_used = metric_used
+                        ,**kwargs)
+
+            if r is None:
+                continue
+
+            _, outputs, ys, best_configs, time_used = r
+
+            if torch.is_tensor(outputs):
+                outputs = outputs.to(outputs.device)
+                ys = ys.to(outputs.device)
+
+            ys = ys.T
+            ds_result[f'{ds_name}_best_configs_at_{eval_position}'] = best_configs
+            ds_result[f'{ds_name}_outputs_at_{eval_position}'] = outputs
+            ds_result[f'{ds_name}_ys_at_{eval_position}'] = ys
+            ds_result[f'{ds_name}_time_at_{eval_position}'] = time_used
+
+            new_metric = torch_nanmean(torch.stack([metric_used(ys[i], outputs[i]) for i in range(ys.shape[0])]))
+
+            if not return_tensor:
+                make_scalar = lambda x: float(x.detach().cpu().numpy()) if (torch.is_tensor(x) and (len(x.shape) == 0)) else x
+                new_metric = make_scalar(new_metric)
+                ds_result = {k: make_scalar(ds_result[k]) for k in ds_result.keys()}
+
+            lib = torch if return_tensor else np
+            if not lib.isnan(new_metric).any():
+                aggregated_metric, num = aggregated_metric + new_metric, num + 1
+
+        overall_result.update(ds_result)
+        if num > 0:
+            aggregated_metric_datasets, num_datasets = (aggregated_metric_datasets + (aggregated_metric / num)), num_datasets + 1
+
+    overall_result['mean_metric'] = aggregated_metric_datasets / num_datasets
+
+    return overall_result
+
+"""
+===============================
+INTERNAL HELPER FUNCTIONS
+===============================
+"""
+
+def check_file_exists(path):
+    """Checks if a pickle file exists. Returns None if not, else returns the unpickled file."""
+    if (os.path.isfile(path)):
+        print(f'loading results from {path}')
+        with open(path, 'rb') as f:
+            return np.load(f, allow_pickle=True).tolist()
+    return None
+
+def generate_valid_split(X, y, bptt, eval_position, split_number=1):
+    """Generates a deteministic train-(test/valid) split. Both splits must contain the same classes and all classes in
+    the entire datasets. If no such split can be sampled in 7 passes, returns None.
+
+    :param X: torch tensor, feature values
+    :param y: torch tensor, class values
+    :param bptt: Number of samples in train + test
+    :param eval_position: Number of samples in train, i.e. from which index values are in test
+    :param split_number: The split id
+    :return:
+    """
+    done, seed = False, 13
+
+    torch.manual_seed(split_number)
+    perm = torch.randperm(X.shape[0]) if split_number > 1 else torch.arange(0, X.shape[0])
+    X, y = X[perm], y[perm]
+
+    while not done:
+        if seed > 20:
+            return None, None # No split could be generated in 7 passes, return None
+        random.seed(seed)
+        i = random.randint(0, len(X) - bptt) if len(X) - bptt > 0 else 0
+        y_ = y[i:i + bptt]
+
+        # Checks if all classes from dataset are contained and classes in train and test are equal (contain same
+        # classes) and
+        done = len(torch.unique(y_)) == len(torch.unique(y))
+        done = done and torch.all(torch.unique(y_) == torch.unique(y))
+        done = done and len(torch.unique(y_[:eval_position])) == len(torch.unique(y_[eval_position:]))
+        done = done and torch.all(torch.unique(y_[:eval_position]) == torch.unique(y_[eval_position:]))
+        seed = seed + 1
+
+    eval_xs = torch.stack([X[i:i + bptt].clone()], 1)
+    eval_ys = torch.stack([y[i:i + bptt].clone()], 1)
+
+    return eval_xs, eval_ys
+
+
+def evaluate_position(X, y, categorical_feats, model, bptt
+                      , eval_position, overwrite, save, base_path, path_interfix, method, ds_name, fetch_only=False
+                      , max_time=300, split_number=1
+                      , per_step_normalization=False, **kwargs):
+    """
+    Evaluates a dataset with a 'bptt' number of training samples.
+
+    :param X: Dataset X
+    :param y: Dataset labels
+    :param categorical_feats: Indices of categorical features.
+    :param model: Model function
+    :param bptt: Sequence length.
+    :param eval_position: Number of training samples.
+    :param overwrite: Wheater to ove
+    :param overwrite: If True, results on disk are overwritten.
+    :param save:
+    :param path_interfix: Used for constructing path to write on disk.
+    :param method: Model name.
+    :param ds_name: Datset name.
+    :param fetch_only: Wheater to calculate or only fetch results.
+    :param per_step_normalization:
+    :param kwargs:
+    :return:
+    """
+
+    if save:
+        path = os.path.join(base_path, f'results/tabular/{path_interfix}/results_{method}_{ds_name}_{eval_position}_{bptt}_{split_number}.npy')
+        #log_path =
+
+    ## Load results if on disk
+    if not overwrite:
+        result = check_file_exists(path)
+        if result is not None:
+            if not fetch_only:
+                print(f'Loaded saved result for {path}')
+            return result
+        elif fetch_only:
+            print(f'Could not load saved result for {path}')
+            return None
+
+    ## Generate data splits
+    eval_xs, eval_ys = generate_valid_split(X, y, bptt, eval_position, split_number=split_number)
+    if eval_xs is None:
+        return None
+        print(f"No dataset could be generated {ds_name} {bptt}")
+
+    eval_ys = (eval_ys > torch.unique(eval_ys).unsqueeze(0)).sum(axis=1).unsqueeze(-1)
+
+    start_time = time.time()
+
+    if isinstance(model, nn.Module): # Two separate predict interfaces for transformer and baselines
+        outputs, best_configs = transformer_predict(model, eval_xs, eval_ys, eval_position, categorical_feats=categorical_feats, **kwargs), None
+    else:
+        _, outputs, best_configs = baseline_predict(model, eval_xs, eval_ys, categorical_feats
+                                                    , eval_pos=eval_position
+                                                    , max_time=max_time, **kwargs)
+
+    eval_ys = eval_ys[eval_position:]
+    if outputs is None:
+        return None
+
+    if torch.is_tensor(outputs): # Transfers data to cpu for saving
+        outputs = outputs.cpu()
+        eval_ys = eval_ys.cpu()
+
+    ds_result = None, outputs, eval_ys, best_configs, time.time() - start_time
+
+    if save:
+        with open(path, 'wb') as f:
+            np.save(f, ds_result)
+            print(f'saved results to {path}')
+
+    return ds_result

TabPFN/scripts/tabular_metrics.py

@@ -0,0 +1,181 @@
+"""
+===============================
+Metrics calculation
+===============================
+Includes a few metric as well as functions composing metrics on results files.
+
+"""
+
+
+
+import numpy as np
+import torch
+from sklearn.metrics import roc_auc_score, accuracy_score, balanced_accuracy_score, average_precision_score
+from scipy.stats import rankdata
+import pandas as pd
+
+"""
+===============================
+Metrics calculation
+===============================
+"""
+def auc_metric(target, pred, multi_class='ovo', numpy=False):
+    lib = np if numpy else torch
+    try:
+        if not numpy:
+            target = torch.tensor(target) if not torch.is_tensor(target) else target
+            pred = torch.tensor(pred) if not torch.is_tensor(pred) else pred
+        if len(lib.unique(target)) > 2:
+            if not numpy:
+                return torch.tensor(roc_auc_score(target, pred, multi_class=multi_class))
+            return roc_auc_score(target, pred, multi_class=multi_class)
+        else:
+            if len(pred.shape) == 2:
+                pred = pred[:, 1]
+            if not numpy:
+                return torch.tensor(roc_auc_score(target, pred))
+            return roc_auc_score(target, pred)
+    except ValueError as e:
+        print(e)
+        return np.nan
+
+def accuracy_metric(target, pred):
+    target = torch.tensor(target) if not torch.is_tensor(target) else target
+    pred = torch.tensor(pred) if not torch.is_tensor(pred) else pred
+    if len(torch.unique(target)) > 2:
+        return torch.tensor(accuracy_score(target, torch.argmax(pred, -1)))
+    else:
+        return torch.tensor(accuracy_score(target, pred[:, 1] > 0.5))
+
+def average_precision_metric(target, pred):
+    target = torch.tensor(target) if not torch.is_tensor(target) else target
+    pred = torch.tensor(pred) if not torch.is_tensor(pred) else pred
+    if len(torch.unique(target)) > 2:
+        return torch.tensor(average_precision_score(target, torch.argmax(pred, -1)))
+    else:
+        return torch.tensor(average_precision_score(target, pred[:, 1] > 0.5))
+
+def balanced_accuracy_metric(target, pred):
+    target = torch.tensor(target) if not torch.is_tensor(target) else target
+    pred = torch.tensor(pred) if not torch.is_tensor(pred) else pred
+    if len(torch.unique(target)) > 2:
+        return torch.tensor(balanced_accuracy_score(target, torch.argmax(pred, -1)))
+    else:
+        return torch.tensor(balanced_accuracy_score(target, pred[:, 1] > 0.5))
+
+def cross_entropy(target, pred):
+    target = torch.tensor(target) if not torch.is_tensor(target) else target
+    pred = torch.tensor(pred) if not torch.is_tensor(pred) else pred
+    if len(torch.unique(target)) > 2:
+        ce = torch.nn.CrossEntropyLoss()
+        return ce(pred.float(), target.long())
+    else:
+        bce = torch.nn.BCELoss()
+        return bce(pred[:, 1].float(), target.float())
+
+def time_metric():
+    """
+    Dummy function, will just be used as a handler.
+    """
+    pass
+
+def count_metric(x, y):
+    """
+    Dummy function, returns one count per dataset.
+    """
+    return 1
+
+"""
+===============================
+Metrics composition
+===============================
+"""
+def calculate_score_per_method(metric, name:str, global_results:dict, ds:list, eval_positions:list, aggregator:str='mean'):
+    """
+    Calculates the metric given by 'metric' and saves it under 'name' in the 'global_results'
+
+    :param metric: Metric function
+    :param name: Name of metric in 'global_results'
+    :param global_results: Dicrtonary containing the results for current method for a collection of datasets
+    :param ds: Dataset to calculate metrics on, a list of dataset properties
+    :param eval_positions: List of positions to calculate metrics on
+    :param aggregator: Specifies way to aggregate results across evaluation positions
+    :return:
+    """
+    aggregator_f = np.nanmean if aggregator == 'mean' else np.nansum
+    for pos in eval_positions:
+        valid_positions = 0
+        for d in ds:
+            if f'{d[0]}_outputs_at_{pos}' in global_results:
+                preds = global_results[f'{d[0]}_outputs_at_{pos}']
+                y = global_results[f'{d[0]}_ys_at_{pos}']
+
+                preds, y = preds.detach().cpu().numpy() if torch.is_tensor(
+                    preds) else preds, y.detach().cpu().numpy() if torch.is_tensor(y) else y
+
+                try:
+                    if metric == time_metric:
+                        global_results[f'{d[0]}_{name}_at_{pos}'] = global_results[f'{d[0]}_time_at_{pos}']
+                        valid_positions = valid_positions + 1
+                    else:
+                        global_results[f'{d[0]}_{name}_at_{pos}'] = aggregator_f(
+                            [metric(y[split], preds[split]) for split in range(y.shape[0])])
+                        valid_positions = valid_positions + 1
+                except Exception as err:
+                    print(f'Error calculating metric with {err}, {type(err)} at {d[0]} {pos} {name}')
+                    global_results[f'{d[0]}_{name}_at_{pos}'] = np.nan
+            else:
+                global_results[f'{d[0]}_{name}_at_{pos}'] = np.nan
+
+        if valid_positions > 0:
+            global_results[f'{aggregator}_{name}_at_{pos}'] = aggregator_f([global_results[f'{d[0]}_{name}_at_{pos}'] for d in ds])
+        else:
+            global_results[f'{aggregator}_{name}_at_{pos}'] = np.nan
+
+    for d in ds:
+        metrics = [global_results[f'{d[0]}_{name}_at_{pos}'] for pos in eval_positions]
+        metrics = [m for m in metrics if not np.isnan(m)]
+        global_results[f'{d[0]}_{aggregator}_{name}'] = aggregator_f(metrics) if len(metrics) > 0 else np.nan
+
+    metrics = [global_results[f'{aggregator}_{name}_at_{pos}'] for pos in eval_positions]
+    metrics = [m for m in metrics if not np.isnan(m)]
+    global_results[f'{aggregator}_{name}'] = aggregator_f(metrics) if len(metrics) > 0 else np.nan
+
+
+def calculate_score(metric, name, global_results, ds, eval_positions, aggregator='mean', limit_to=''):
+    """
+    Calls calculate_metrics_by_method with a range of methods. See arguments of that method.
+    :param limit_to: This method will not get metric calculations.
+    """
+    for m in global_results:
+        if limit_to not in m:
+            continue
+        calculate_score_per_method(metric, name, global_results[m], ds, eval_positions, aggregator=aggregator)
+
+
+def make_metric_matrix(global_results, methods, pos, name, ds):
+    result = []
+    for m in global_results:
+        result += [[global_results[m][d[0] + '_' + name + '_at_' + str(pos)] for d in ds]]
+    result = np.array(result)
+    result = pd.DataFrame(result.T, index=[d[0] for d in ds], columns=[k[:-8] for k in list(global_results.keys())])
+
+    matrix_means, matrix_stds = [], []
+
+    for method in methods:
+        matrix_means += [result.iloc[:, [(method) in c for c in result.columns]].mean(axis=1)]
+        matrix_stds += [result.iloc[:, [(method) in c for c in result.columns]].std(axis=1)]
+
+    matrix_means = pd.DataFrame(matrix_means, index=methods).T
+    matrix_stds = pd.DataFrame(matrix_stds, index=methods).T
+
+    return matrix_means, matrix_stds
+
+
+def make_ranks_and_wins_table(matrix):
+    for dss in matrix.T:
+        matrix.loc[dss] = rankdata(-matrix.round(3).loc[dss])
+    ranks_acc = matrix.mean()
+    wins_acc = (matrix == 1).sum()
+
+    return ranks_acc, wins_acc

TabPFN/scripts/transformer_prediction_interface.py

@@ -0,0 +1,357 @@
+import torch
+import random
+
+from torch.utils.checkpoint import checkpoint
+
+from utils import normalize_data, to_ranking_low_mem, remove_outliers
+from priors.utils import normalize_by_used_features_f
+from utils import NOP
+
+from sklearn.preprocessing import PowerTransformer, QuantileTransformer, RobustScaler
+
+from notebook_utils import CustomUnpickler
+
+import numpy as np
+from sklearn.base import BaseEstimator, ClassifierMixin
+from sklearn.utils.validation import check_X_y, check_array, check_is_fitted
+from sklearn.utils.multiclass import check_classification_targets
+from sklearn.utils import column_or_1d
+from pathlib import Path
+from model_builder import load_model
+import os
+
+def load_model_workflow(i, e, add_name, base_path, device='cpu', eval_addition=''):
+    """
+    Workflow for loading a model and setting appropriate parameters for diffable hparam tuning.
+
+    :param i:
+    :param e:
+    :param eval_positions_valid:
+    :param add_name:
+    :param base_path:
+    :param device:
+    :param eval_addition:
+    :return:
+    """
+    def check_file(e):
+        model_file = f'models_diff/prior_diff_real_checkpoint{add_name}_n_{i}_epoch_{e}.cpkt'
+        model_path = os.path.join(base_path, model_file)
+        # print('Evaluate ', model_path)
+        results_file = os.path.join(base_path,
+                                    f'models_diff/prior_diff_real_results{add_name}_n_{i}_epoch_{e}_{eval_addition}.pkl')
+        if not Path(model_path).is_file():  # or Path(results_file).is_file():
+            return None, None, None
+        return model_file, model_path, results_file
+
+    model_file = None
+    if e == -1:
+        for e_ in range(100, -1, -1):
+            model_file_, model_path_, results_file_ = check_file(e_)
+            if model_file_ is not None:
+                e = e_
+                model_file, model_path, results_file = model_file_, model_path_, results_file_
+                break
+    else:
+        model_file, model_path, results_file = check_file(e)
+
+    if model_file is None:
+        print('No checkpoint found')
+        return None
+
+    print(f'Loading {model_file}')
+
+    model, c = load_model(base_path, model_file, device, eval_positions=[], verbose=False)
+
+    return model, c, results_file
+
+
+class TabPFNClassifier(BaseEstimator, ClassifierMixin):
+
+    def __init__(self, device='cpu', base_path='.'):
+        # Model file specification (Model name, Epoch)
+        model_string = ''
+        i, e = '8x_lr0.0003', -1
+
+        # File which contains result of hyperparameter tuning run: style (i.e. hyperparameters) and a dataframe with results.
+        style_file = 'prior_tuning_result.pkl'
+
+        model, c, results_file = load_model_workflow(i, e, add_name=model_string, base_path=base_path, device=device,
+                                                     eval_addition='')
+        style, temperature = self.load_result_minimal(style_file, i, e, base_path=base_path)
+
+        self.device = device
+        self.base_path = base_path
+        self.model = model
+        self.c = c
+        self.style = style
+        self.temperature = temperature
+
+        self.max_num_features = self.c['num_features']
+        self.max_num_classes = self.c['max_num_classes']
+
+    def load_result_minimal(self, path, i, e, base_path='.'):
+        with open(os.path.join(base_path,path), 'rb') as output:
+            _, _, _, style, temperature, optimization_route = CustomUnpickler(output).load()
+
+            return style, temperature
+
+    def fit(self, X, y):
+        # Check that X and y have correct shape
+        X, y = check_X_y(X, y)
+        y = self._validate_targets(y)
+
+        self.X_ = X
+        self.y_ = y
+
+        if X.shape[1] > self.max_num_features:
+            raise ValueError("The number of features for this classifier is restricted to ", self.max_num_features)
+        if len(np.unique(y)) > self.max_num_classes:
+            raise ValueError("The number of classes for this classifier is restricted to ", self.max_num_classes)
+
+        # Return the classifier
+        return self
+
+    def _validate_targets(self, y):
+        y_ = column_or_1d(y, warn=True)
+        check_classification_targets(y)
+        cls, y = np.unique(y_, return_inverse=True)
+        if len(cls) < 2:
+            raise ValueError(
+                "The number of classes has to be greater than one; got %d class"
+                % len(cls)
+            )
+
+        self.classes_ = cls
+
+        return np.asarray(y, dtype=np.float64, order="C")
+
+    def predict_proba(self, X):
+        # Check is fit had been called
+        check_is_fitted(self)
+
+        # Input validation
+        X = check_array(X)
+
+        X_full = np.concatenate([self.X_, X], axis=0)
+        X_full = torch.tensor(X_full, device=self.device).float().unsqueeze(1)
+        y_full = np.concatenate([self.y_, self.y_[0] + np.zeros_like(X[:, 0])], axis=0)
+        y_full = torch.tensor(y_full, device=self.device).float().unsqueeze(1)
+
+        eval_pos = self.X_.shape[0]
+
+        prediction = transformer_predict(self.model[2], X_full, y_full, eval_pos,
+                                         device=self.device,
+                                         style=self.style,
+                                         inference_mode=True,
+                                         N_ensemble_configurations=10,
+                                         softmax_temperature=self.temperature
+                                         , **get_params_from_config(self.c))
+        prediction_ = prediction.squeeze(0)
+
+        return prediction_.detach().cpu().numpy()
+
+    def predict(self, X, return_winning_probability=False):
+        p = self.predict_proba(X)
+        y = np.argmax(self.predict_proba(X), axis=-1)
+        y = self.classes_.take(np.asarray(y, dtype=np.intp))
+        if return_winning_probability:
+            return y, p.max(axis=-1)
+        return y
+
+def transformer_predict(model, eval_xs, eval_ys, eval_position,
+                        device='cpu',
+                        max_features=100,
+                        style=None,
+                        inference_mode=False,
+                        num_classes=2,
+                        extend_features=True,
+                        normalize_to_ranking=False,
+                        softmax_temperature=0.0,
+                        multiclass_decoder='permutation',
+                        preprocess_transform='mix',
+                        categorical_feats=[],
+                        feature_shift_decoder=True,
+                        N_ensemble_configurations=10,
+                        average_logits=True,
+                        normalize_with_sqrt=False, **kwargs):
+    """
+
+    :param model:
+    :param eval_xs:
+    :param eval_ys: should be classes that are 0-indexed and every class until num_classes-1 is present
+    :param eval_position:
+    :param rescale_features:
+    :param device:
+    :param max_features:
+    :param style:
+    :param inference_mode:
+    :param num_classes:
+    :param extend_features:
+    :param normalize_to_ranking:
+    :param softmax_temperature:
+    :param multiclass_decoder:
+    :param preprocess_transform:
+    :param categorical_feats:
+    :param feature_shift_decoder:
+    :param N_ensemble_configurations:
+    :param average_logits:
+    :param normalize_with_sqrt:
+    :param metric_used:
+    :return:
+    """
+    num_classes = len(torch.unique(eval_ys))
+
+    def predict(eval_xs, eval_ys, used_style, softmax_temperature, return_logits):
+        # Initialize results array size S, B, Classes
+
+        inference_mode_call = torch.inference_mode() if inference_mode else NOP()
+        with inference_mode_call:
+            output = model(
+                    (used_style.repeat(eval_xs.shape[1], 1) if used_style is not None else None, eval_xs, eval_ys.float()),
+                    single_eval_pos=eval_position)[:, :, 0:num_classes]
+
+            output = output[:, :, 0:num_classes] / torch.exp(softmax_temperature)
+            if not return_logits:
+                output = torch.nn.functional.softmax(output, dim=-1)
+            #else:
+            #    output[:, :, 1] = model((style.repeat(eval_xs.shape[1], 1) if style is not None else None, eval_xs, eval_ys.float()),
+            #               single_eval_pos=eval_position)
+
+            #    output[:, :, 1] = torch.sigmoid(output[:, :, 1]).squeeze(-1)
+            #    output[:, :, 0] = 1 - output[:, :, 1]
+
+        #print('RESULTS', eval_ys.shape, torch.unique(eval_ys, return_counts=True), output.mean(axis=0))
+
+        return output
+
+    def preprocess_input(eval_xs, preprocess_transform):
+        import warnings
+
+        if eval_xs.shape[1] > 1:
+            raise Exception("Transforms only allow one batch dim - TODO")
+        if preprocess_transform != 'none':
+            if preprocess_transform == 'power' or preprocess_transform == 'power_all':
+                pt = PowerTransformer(standardize=True)
+            elif preprocess_transform == 'quantile' or preprocess_transform == 'quantile_all':
+                pt = QuantileTransformer(output_distribution='normal')
+            elif preprocess_transform == 'robust' or preprocess_transform == 'robust_all':
+                pt = RobustScaler(unit_variance=True)
+
+        # eval_xs, eval_ys = normalize_data(eval_xs), normalize_data(eval_ys)
+        eval_xs = normalize_data(eval_xs)
+
+        # Removing empty features
+        eval_xs = eval_xs[:, 0, :].cpu().numpy()
+        sel = [len(np.unique(eval_xs[0:eval_ys.shape[0], col])) > 1 for col in range(eval_xs.shape[1])]
+        eval_xs = np.array(eval_xs[:, sel])
+
+        warnings.simplefilter('error')
+        if preprocess_transform != 'none':
+            feats = set(range(eval_xs.shape[1])) if 'all' in preprocess_transform else set(
+                range(eval_xs.shape[1])) - set(categorical_feats)
+            for col in feats:
+                try:
+                    pt.fit(eval_xs[0:eval_ys.shape[0], col:col + 1])
+                    trans = pt.transform(eval_xs[:, col:col + 1])
+                    # print(scipy.stats.spearmanr(trans[~np.isnan(eval_xs[:, col:col+1])], eval_xs[:, col:col+1][~np.isnan(eval_xs[:, col:col+1])]))
+                    eval_xs[:, col:col + 1] = trans
+                except:
+                    pass
+        warnings.simplefilter('default')
+
+        eval_xs = torch.tensor(eval_xs).float().unsqueeze(1).to(device)
+
+        # eval_xs = normalize_data(eval_xs)
+
+        # TODO: Cautian there is information leakage when to_ranking is used, we should not use it
+        eval_xs = remove_outliers(eval_xs) if not normalize_to_ranking else normalize_data(to_ranking_low_mem(eval_xs))
+
+        # Rescale X
+        eval_xs = normalize_by_used_features_f(eval_xs, eval_xs.shape[-1], max_features,
+                                               normalize_with_sqrt=normalize_with_sqrt)
+        return eval_xs.detach()
+
+    eval_xs, eval_ys = eval_xs.to(device), eval_ys.to(device)
+    eval_ys = eval_ys[:eval_position]
+
+    model.to(device)
+    style = style.to(device)
+
+    model.eval()
+
+    import itertools
+    style = style.unsqueeze(0) if len(style.shape) == 1 else style
+    num_styles = style.shape[0]
+    styles_configurations = range(0, num_styles)
+    preprocess_transform_configurations = [preprocess_transform if i % 2 == 0 else 'none' for i in range(0, num_styles)]
+    if preprocess_transform == 'mix':
+        def get_preprocess(i):
+            if i == 0:
+                return 'power_all'
+            if i == 1:
+                return 'robust_all'
+            if i == 2:
+                return 'none'
+        preprocess_transform_configurations = [get_preprocess(i) for i in range(0, num_styles)]
+    styles_configurations = zip(styles_configurations, preprocess_transform_configurations)
+
+    feature_shift_configurations = range(0, eval_xs.shape[2]) if feature_shift_decoder else [0]
+    class_shift_configurations = range(0, len(torch.unique(eval_ys))) if multiclass_decoder == 'permutation' else [0]
+
+    ensemble_configurations = list(itertools.product(styles_configurations, feature_shift_configurations, class_shift_configurations))
+    random.shuffle(ensemble_configurations)
+    ensemble_configurations = ensemble_configurations[0:N_ensemble_configurations]
+
+    output = None
+
+    eval_xs_transformed = {}
+    for ensemble_configuration in ensemble_configurations:
+        (styles_configuration, preprocess_transform_configuration), feature_shift_configuration, class_shift_configuration = ensemble_configuration
+
+        style_ = style[styles_configuration:styles_configuration+1, :]
+        softmax_temperature_ = softmax_temperature[styles_configuration]
+
+        eval_xs_, eval_ys_ = eval_xs.clone(), eval_ys.clone()
+
+        if preprocess_transform_configuration in eval_xs_transformed:
+            eval_xs_ = eval_xs_transformed['preprocess_transform_configuration'].clone()
+        else:
+            eval_xs_ = preprocess_input(eval_xs_, preprocess_transform=preprocess_transform_configuration)
+            eval_xs_transformed['preprocess_transform_configuration'] = eval_xs_
+
+        eval_ys_ = ((eval_ys_ + class_shift_configuration) % num_classes).float()
+
+        eval_xs_ = torch.cat([eval_xs_[..., feature_shift_configuration:],eval_xs_[..., :feature_shift_configuration]],dim=-1)
+
+        # Extend X
+        if extend_features:
+            eval_xs_ = torch.cat(
+                [eval_xs_,
+                 torch.zeros((eval_xs_.shape[0], eval_xs_.shape[1], max_features - eval_xs_.shape[2])).to(device)], -1)
+
+        #preprocess_transform_ = preprocess_transform if styles_configuration % 2 == 0 else 'none'
+        import warnings
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", message="None of the inputs have requires_grad=True. Gradients will be None")
+            output_ = checkpoint(predict, eval_xs_, eval_ys_, style_, softmax_temperature_, True)
+            output_ = torch.cat([output_[..., class_shift_configuration:],output_[..., :class_shift_configuration]],dim=-1)
+
+        #output_ = predict(eval_xs, eval_ys, style_, preprocess_transform_)
+        if not average_logits:
+            output_ = torch.nn.functional.softmax(output_, dim=-1)
+        output = output_ if output is None else output + output_
+
+    output = output / len(ensemble_configurations)
+    if average_logits:
+        output = torch.nn.functional.softmax(output, dim=-1)
+
+    output = torch.transpose(output, 0, 1)
+
+    return output
+
+def get_params_from_config(c):
+    return {'max_features': c['num_features']
+        , 'rescale_features': c["normalize_by_used_features"]
+        , 'normalize_to_ranking': c["normalize_to_ranking"]
+        , 'normalize_with_sqrt': c.get("normalize_with_sqrt", False)
+            }

TabPFN/tabular_evaluation.py

@@ -0,0 +1,283 @@
+import time
+import os
+from pathlib import Path
+
+from tqdm import tqdm
+import random
+import numpy as np
+
+from torch import nn
+
+from utils import torch_nanmean
+from datasets import *
+from model_builder import load_model
+from scripts.tabular_baselines import get_scoring_string
+from scripts import tabular_metrics
+from scripts.transformer_prediction_interface import *
+from scripts.baseline_prediction_interface import *
+"""
+===============================
+PUBLIC FUNCTIONS FOR EVALUATION
+===============================
+"""
+
+
+def eval_model(i, e, valid_datasets, test_datasets, eval_positions, bptt, add_name, base_path, device='cpu', eval_addition='', **kwargs):
+    metrics_test, config_sample, model_path = eval_model_on_ds(i, e, test_datasets, eval_positions, bptt, add_name, base_path, device=device, eval_addition=eval_addition, **kwargs)
+    metrics_valid, _, _ = eval_model_on_ds(i, e, valid_datasets, eval_positions, bptt, add_name, base_path, device=device, eval_addition=eval_addition, **kwargs)
+    return {'mean_auc_test': metrics_test['mean_roc_at_1000'], 'mean_auc_valid': metrics_valid['mean_roc_at_1000'], 'mean_ce_test': metrics_test['mean_ce_at_1000'], 'mean_ce_valid': metrics_valid['mean_ce_at_1000'], 'config_sample': config_sample, 'model_path': model_path}
+
+def eval_model_on_ds(i, e, valid_datasets, eval_positions, bptt, add_name, base_path, device='cpu', eval_addition='', **kwargs):
+
+    # How to use: evaluate_without_fitting(i,0,valid_datasets, [1024], 100000, add_name=model_string, base_path=base_path,)
+    def check_file(e):
+        model_file = f'models_diff/prior_diff_real_checkpoint{add_name}_n_{i}_epoch_{e}.cpkt'
+        model_path = os.path.join(base_path, model_file)
+        # print('Evaluate ', model_path)
+        results_file = os.path.join(base_path,
+                                    f'models_diff/prior_diff_real_results{add_name}_n_{i}_epoch_{e}_{eval_addition}.pkl')
+        if not Path(model_path).is_file():  # or Path(results_file).is_file():
+            # print('checkpoint exists: ', Path(model_file).is_file(), ', results are written:', Path(results_file).is_file())
+            return None, None, None
+        return model_file, model_path, results_file
+
+    if e == -1: # use last checkpoint, if e == -1
+        for e_ in range(100, -1, -1):
+            model_file_, model_path_, results_file_ = check_file(e_)
+            if model_file_ is not None:
+                e = e_
+                model_file, model_path, results_file = model_file_, model_path_, results_file_
+                break
+    else:
+        model_file, model_path, results_file = check_file(e)
+
+    model, config_sample = load_model(base_path, model_file, device, None, verbose=False)
+
+    params = {'max_features': config_sample['num_features']
+        , 'rescale_features': config_sample["normalize_by_used_features"]
+        , 'normalize_to_ranking': config_sample["normalize_to_ranking"]
+        , 'normalize_with_sqrt': config_sample.get("normalize_with_sqrt", False)
+              }
+    metrics_valid = evaluate(datasets=valid_datasets, model=model[2], method='transformer', device=device, overwrite=True,
+                             extend_features=True
+                             # just removed the style keyword but transformer is trained with style, just empty
+                             , save=False
+                             , metric_used=tabular_metrics.cross_entropy
+                             , return_tensor=True
+                             , verbose=False
+                             , eval_positions=eval_positions
+                             , bptt=bptt
+                             , base_path=None
+                             , inference_mode=True
+                             , **params
+                             , **kwargs)
+
+    tabular_metrics.calculate_score_per_method(tabular_metrics.auc_metric, 'roc', metrics_valid, valid_datasets, eval_positions)
+    tabular_metrics.calculate_score_per_method(tabular_metrics.cross_entropy, 'ce', metrics_valid, valid_datasets, eval_positions)
+
+    return metrics_valid, config_sample, model_path
+
+
+def evaluate(datasets, bptt, eval_positions, metric_used, model
+             , verbose=False
+             , return_tensor=False
+             , **kwargs):
+    """
+    Evaluates a list of datasets for a model function.
+
+    :param datasets: List of datasets
+    :param bptt: maximum sequence length
+    :param eval_positions: List of positions where to evaluate models
+    :param verbose: If True, is verbose.
+    :param metric_used: Which metric is optimized for.
+    :param return_tensor: Wheater to return results as a pytorch.tensor or numpy, this is only relevant for transformer.
+    :param kwargs:
+    :return:
+    """
+    overall_result = {'metric_used': get_scoring_string(metric_used)
+                      , 'bptt': bptt
+                      , 'eval_positions': eval_positions}
+
+    aggregated_metric_datasets, num_datasets = torch.tensor(0.0), 0
+
+    # For each dataset
+    for [ds_name, X, y, categorical_feats, _, _] in tqdm.tqdm(datasets, desc='Iterate over datasets') if verbose else datasets:
+        dataset_bptt = min(len(X), bptt)
+        # if verbose and dataset_bptt < bptt:
+        #    print(f'Dataset too small for given sequence length, reducing to {len(X)} ({bptt})')
+
+        aggregated_metric, num = torch.tensor(0.0), 0
+        ds_result = {}
+
+        for eval_position in (eval_positions if verbose else eval_positions):
+            eval_position_real = int(dataset_bptt * 0.5) if 2 * eval_position > dataset_bptt else eval_position
+            eval_position_bptt = int(eval_position_real * 2.0)
+
+            r = evaluate_position(X, y, model=model
+                        , num_classes=len(torch.unique(y))
+                        , categorical_feats = categorical_feats
+                        , bptt = eval_position_bptt
+                        , ds_name=ds_name
+                        , eval_position = eval_position_real
+                        , metric_used = metric_used
+                        ,**kwargs)
+
+            if r is None:
+                continue
+
+            _, outputs, ys, best_configs, time_used = r
+
+            if torch.is_tensor(outputs):
+                outputs = outputs.to(outputs.device)
+                ys = ys.to(outputs.device)
+
+            ys = ys.T
+            ds_result[f'{ds_name}_best_configs_at_{eval_position}'] = best_configs
+            ds_result[f'{ds_name}_outputs_at_{eval_position}'] = outputs
+            ds_result[f'{ds_name}_ys_at_{eval_position}'] = ys
+            ds_result[f'{ds_name}_time_at_{eval_position}'] = time_used
+
+            new_metric = torch_nanmean(torch.stack([metric_used(ys[i], outputs[i]) for i in range(ys.shape[0])]))
+
+            if not return_tensor:
+                make_scalar = lambda x: float(x.detach().cpu().numpy()) if (torch.is_tensor(x) and (len(x.shape) == 0)) else x
+                new_metric = make_scalar(new_metric)
+                ds_result = {k: make_scalar(ds_result[k]) for k in ds_result.keys()}
+
+            lib = torch if return_tensor else np
+            if not lib.isnan(new_metric).any():
+                aggregated_metric, num = aggregated_metric + new_metric, num + 1
+
+        overall_result.update(ds_result)
+        if num > 0:
+            aggregated_metric_datasets, num_datasets = (aggregated_metric_datasets + (aggregated_metric / num)), num_datasets + 1
+
+    overall_result['mean_metric'] = aggregated_metric_datasets / num_datasets
+
+    return overall_result
+
+"""
+===============================
+INTERNAL HELPER FUNCTIONS
+===============================
+"""
+
+def check_file_exists(path):
+    """Checks if a pickle file exists. Returns None if not, else returns the unpickled file."""
+    if (os.path.isfile(path)):
+        print(f'loading results from {path}')
+        with open(path, 'rb') as f:
+            return np.load(f, allow_pickle=True).tolist()
+    return None
+
+def generate_valid_split(X, y, bptt, eval_position, split_number=1):
+    """Generates a deteministic train-(test/valid) split. Both splits must contain the same classes and all classes in
+    the entire datasets. If no such split can be sampled in 7 passes, returns None.
+
+    :param X: torch tensor, feature values
+    :param y: torch tensor, class values
+    :param bptt: Number of samples in train + test
+    :param eval_position: Number of samples in train, i.e. from which index values are in test
+    :param split_number: The split id
+    :return:
+    """
+    done, seed = False, 13
+
+    torch.manual_seed(split_number)
+    perm = torch.randperm(X.shape[0]) if split_number > 1 else torch.arange(0, X.shape[0])
+    X, y = X[perm], y[perm]
+
+    while not done:
+        if seed > 20:
+            return None, None # No split could be generated in 7 passes, return None
+        random.seed(seed)
+        i = random.randint(0, len(X) - bptt) if len(X) - bptt > 0 else 0
+        y_ = y[i:i + bptt]
+
+        # Checks if all classes from dataset are contained and classes in train and test are equal (contain same
+        # classes) and
+        done = len(torch.unique(y_)) == len(torch.unique(y))
+        done = done and torch.all(torch.unique(y_) == torch.unique(y))
+        done = done and len(torch.unique(y_[:eval_position])) == len(torch.unique(y_[eval_position:]))
+        done = done and torch.all(torch.unique(y_[:eval_position]) == torch.unique(y_[eval_position:]))
+        seed = seed + 1
+
+    eval_xs = torch.stack([X[i:i + bptt].clone()], 1)
+    eval_ys = torch.stack([y[i:i + bptt].clone()], 1)
+
+    return eval_xs, eval_ys
+
+
+def evaluate_position(X, y, categorical_feats, model, bptt
+                      , eval_position, overwrite, save, base_path, path_interfix, method, ds_name, fetch_only=False
+                      , max_time=300, split_number=1
+                      , per_step_normalization=False, **kwargs):
+    """
+    Evaluates a dataset with a 'bptt' number of training samples.
+
+    :param X: Dataset X
+    :param y: Dataset labels
+    :param categorical_feats: Indices of categorical features.
+    :param model: Model function
+    :param bptt: Sequence length.
+    :param eval_position: Number of training samples.
+    :param overwrite: Wheater to ove
+    :param overwrite: If True, results on disk are overwritten.
+    :param save:
+    :param path_interfix: Used for constructing path to write on disk.
+    :param method: Model name.
+    :param ds_name: Datset name.
+    :param fetch_only: Wheater to calculate or only fetch results.
+    :param per_step_normalization:
+    :param kwargs:
+    :return:
+    """
+
+    if save:
+        path = os.path.join(base_path, f'results/tabular/{path_interfix}/results_{method}_{ds_name}_{eval_position}_{bptt}_{split_number}.npy')
+        #log_path =
+
+    ## Load results if on disk
+    if not overwrite:
+        result = check_file_exists(path)
+        if result is not None:
+            if not fetch_only:
+                print(f'Loaded saved result for {path}')
+            return result
+        elif fetch_only:
+            print(f'Could not load saved result for {path}')
+            return None
+
+    ## Generate data splits
+    eval_xs, eval_ys = generate_valid_split(X, y, bptt, eval_position, split_number=split_number)
+    if eval_xs is None:
+        return None
+        print(f"No dataset could be generated {ds_name} {bptt}")
+
+    eval_ys = (eval_ys > torch.unique(eval_ys).unsqueeze(0)).sum(axis=1).unsqueeze(-1)
+
+    start_time = time.time()
+
+    if isinstance(model, nn.Module): # Two separate predict interfaces for transformer and baselines
+        outputs, best_configs = transformer_predict(model, eval_xs, eval_ys, eval_position, categorical_feats=categorical_feats, **kwargs), None
+    else:
+        _, outputs, best_configs = baseline_predict(model, eval_xs, eval_ys, categorical_feats
+                                                    , eval_pos=eval_position
+                                                    , max_time=max_time, **kwargs)
+
+    eval_ys = eval_ys[eval_position:]
+    if outputs is None:
+        return None
+
+    if torch.is_tensor(outputs): # Transfers data to cpu for saving
+        outputs = outputs.cpu()
+        eval_ys = eval_ys.cpu()
+
+    ds_result = None, outputs, eval_ys, best_configs, time.time() - start_time
+
+    if save:
+        with open(path, 'wb') as f:
+            np.save(f, ds_result)
+            print(f'saved results to {path}')
+
+    return ds_result

TabPFN/train.py

@@ -0,0 +1,386 @@
+import os
+import itertools
+import argparse
+import time
+import datetime
+import yaml
+from contextlib import nullcontext
+
+
+import torch
+from torch import nn
+
+import utils
+from transformer import TransformerModel
+from utils import get_cosine_schedule_with_warmup, get_openai_lr, StoreDictKeyPair, get_weighted_single_eval_pos_sampler, get_uniform_single_eval_pos_sampler
+import priors
+import encoders
+import positional_encodings
+from utils import init_dist
+from torch.cuda.amp import autocast
+
+class Losses():
+    gaussian = nn.GaussianNLLLoss(full=True, reduction='none')
+    mse = nn.MSELoss(reduction='none')
+    ce = lambda weight : nn.CrossEntropyLoss(reduction='none', weight=weight)
+    bce = nn.BCEWithLogitsLoss(reduction='none')
+
+
+def train(priordataloader_class, criterion, encoder_generator, emsize=200, nhid=200, nlayers=6, nhead=2, dropout=0.2,
+          epochs=10, steps_per_epoch=100, batch_size=200, bptt=10, lr=None, weight_decay=0.0, warmup_epochs=10, input_normalization=False,
+          y_encoder_generator=None, pos_encoder_generator=None, decoder=None, extra_prior_kwargs_dict={}, scheduler=get_cosine_schedule_with_warmup,
+          load_weights_from_this_state_dict=None, validation_period=10, single_eval_pos_gen=None, bptt_extra_samples=None, gpu_device='cuda:0',
+          aggregate_k_gradients=1, verbose=True, style_encoder_generator=None, check_is_compatible=True, epoch_callback=None,
+          initializer=None, initialize_with_model=None, train_mixed_precision=False, total_available_time_in_s=None, normalize_labels=True, **model_extra_args
+          ):
+    assert (epochs is None) != (total_available_time_in_s is None)
+    start_of_training = time.time()
+    device = gpu_device if torch.cuda.is_available() else 'cpu:0'
+    print(f'Using {device} device')
+    using_dist, rank, device = init_dist(device)
+    bptt_sampler = (lambda : single_eval_pos_gen() + bptt_extra_samples if callable(single_eval_pos_gen) else single_eval_pos_gen + bptt_extra_samples) if bptt_extra_samples is not None else bptt
+    dl = priordataloader_class(num_steps=steps_per_epoch, batch_size=batch_size, seq_len=bptt_sampler, seq_len_maximum=bptt+(bptt_extra_samples if bptt_extra_samples else 0), device=device, **extra_prior_kwargs_dict)
+    if dl.fuse_x_y:
+        raise Exception("Illegal parameter")
+
+    encoder = encoder_generator(dl.num_features+1 if dl.fuse_x_y else dl.num_features,emsize)
+    style_def = next(iter(dl))[0][0] # This is (style, x, y), target with x and y with batch size
+
+    style_encoder = style_encoder_generator(hyperparameter_definitions=style_def[0], em_size=emsize) if (style_def is not None) else None
+    n_out = dl.num_outputs
+    if isinstance(criterion, nn.GaussianNLLLoss):
+        n_out *= 2
+    elif isinstance(criterion, nn.CrossEntropyLoss):
+        n_out *= criterion.weight.shape[0]
+    model = TransformerModel(encoder, n_out, emsize, nhead, nhid, nlayers, dropout, style_encoder=style_encoder,
+                             y_encoder=y_encoder_generator(dl.num_outputs, emsize), input_normalization=input_normalization,
+                             pos_encoder=(pos_encoder_generator or positional_encodings.NoPositionalEncoding)(emsize, bptt*2),
+                             decoder=decoder, init_method=initializer, **model_extra_args
+                             )
+    model.criterion = criterion
+    if load_weights_from_this_state_dict is not None:
+        model.load_state_dict(load_weights_from_this_state_dict)
+    if initialize_with_model is not None:
+        model.init_from_small_model(initialize_with_model)
+
+    print(f"Using a Transformer with {sum(p.numel() for p in model.parameters())/1000/1000:.{2}f} M parameters")
+
+    try:
+        for (k, v), (k2, v2) in zip(model.state_dict().items(), initialize_with_model.state_dict().items()):
+            print(k, ((v - v2) / v).abs().mean(), v.shape)
+    except Exception:
+        pass
+
+    model.to(device)
+    if using_dist:
+        print("Distributed training")
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank], output_device=rank, broadcast_buffers=False)
+
+
+    # learning rate
+    if lr is None:
+        lr = get_openai_lr(model)
+        print(f"Using OpenAI max lr of {lr}.")
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
+    scheduler = scheduler(optimizer, warmup_epochs, epochs if epochs is not None else 100) # when training for fixed time lr schedule takes 100 steps
+
+    def train_step():
+        model.train()  # Turn on the train mode
+        total_loss = 0.
+        total_positional_losses = 0.
+        total_positional_losses_recorded = 0
+        before_get_batch = time.time()
+        assert len(dl) % aggregate_k_gradients == 0, 'Please set the number of steps per epoch s.t. `aggregate_k_gradients` divides it.'
+        valid_batch_steps = 0.0
+        for batch, (data, targets) in enumerate(dl):
+            if using_dist and not (batch % aggregate_k_gradients == aggregate_k_gradients - 1):
+                cm = model.no_sync()
+                #print(f'p={rank}, no_sync', force=True)
+            else:
+                cm = nullcontext()
+                #print(f'p={rank}, sync', force=True)
+            with cm:
+                time_to_get_batch = time.time() - before_get_batch
+                before_forward = time.time()
+                if bptt_extra_samples is None:
+                    single_eval_pos = single_eval_pos_gen() if callable(single_eval_pos_gen) else single_eval_pos_gen
+                else:
+                    single_eval_pos = targets.shape[0] - bptt_extra_samples
+
+                is_compatible = torch.ones((targets.shape[1])).bool()
+                if check_is_compatible or normalize_labels:
+                    for b in range(targets.shape[1]):
+                        targets_in_train = torch.unique(targets[:single_eval_pos, b], sorted=True)
+                        targets_in_eval = torch.unique(targets[single_eval_pos:, b], sorted=True)
+
+                        if check_is_compatible:
+                            is_compatible[b] = len(targets_in_train) == len(targets_in_eval) and (targets_in_train == targets_in_eval).all()
+                            is_compatible[b] = is_compatible[b] and len(targets_in_train) > 1
+
+                        # Set targets to range starting from 0 (e.g. targets 0, 2, 5, 2 will be converted to 0, 1, 2, 1)
+                        if normalize_labels:
+                            targets[:, b] = (targets[:, b] > torch.unique(targets[:, b]).unsqueeze(1)).sum(axis=0).unsqueeze(0)
+                valid_batch_steps += is_compatible.float().mean()
+                is_compatible = is_compatible.to(device)
+                #if using_dist and check_is_compatible:
+                #    print('step share before reduce',curr_step_share, force=True)
+                #    curr_step_share = curr_step_share.to(device)
+                #    torch.distributed.all_reduce_multigpu([curr_step_share], op=torch.distributed.ReduceOp.SUM)
+                #    curr_step_share = curr_step_share.cpu() / torch.distributed.get_world_size()
+                #    print('step share after reduce',curr_step_share, torch.distributed.get_world_size(), force=True)
+
+                # If style is set to None, it should not be transferred to device
+                output = model(tuple(e.to(device) if torch.is_tensor(e) else e for e in data) if isinstance(data, tuple) else data.to(device)
+                               , single_eval_pos=single_eval_pos)
+
+                forward_time = time.time() - before_forward
+
+                #output, targets = output[:, is_compatible], targets[:, is_compatible]
+
+                if single_eval_pos is not None:
+                    targets = targets[single_eval_pos:]
+                if isinstance(criterion, nn.GaussianNLLLoss):
+                    assert output.shape[-1] == 2, \
+                        'need to write a little bit of code to handle multiple regression targets at once'
+
+                    mean_pred = output[..., 0]
+                    var_pred = output[..., 1].abs()
+                    losses = criterion(mean_pred.flatten(), targets.to(device).flatten(), var=var_pred.flatten())
+                elif isinstance(criterion, (nn.MSELoss, nn.BCEWithLogitsLoss)):
+                    losses = criterion(output.flatten(), targets.to(device).flatten())
+                elif isinstance(criterion, (nn.CrossEntropyLoss)):
+                    #print(n_out, targets.min(), targets.max(), force=True)
+                    losses = criterion(output.reshape(-1, n_out), targets.to(device).long().flatten())
+                else:
+                    losses = criterion(output.reshape(-1, n_out), targets.to(device).flatten())
+                losses = losses.view(*output.shape[0:2])
+                loss = losses.mean(0) @ is_compatible.float() / losses.shape[1]
+                #loss = torch_nanmean(losses, axis=[0, 1]) * is_compatible.float().mean()
+                # not sure whether we can go without the nan checks.
+
+                loss.backward()
+
+                if ((batch % aggregate_k_gradients == aggregate_k_gradients - 1) and (not check_is_compatible or using_dist))\
+                        or (valid_batch_steps >= aggregate_k_gradients and (check_is_compatible and not using_dist)):
+                    with torch.no_grad():
+                        for p in model.parameters():
+                            if p.grad is not None:
+                                p.grad.div_(valid_batch_steps)
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
+                    try:
+                        optimizer.step()
+                    except:
+                        print("Invalid optimization step encountered")
+                    optimizer.zero_grad()
+                    valid_batch_steps = 0.0
+
+                step_time = time.time() - before_forward
+
+                if not torch.isnan(loss):
+                    total_loss += loss.item()
+                    total_positional_losses += losses.mean(1).cpu().detach() if single_eval_pos is None else \
+                        nn.functional.one_hot(torch.tensor(single_eval_pos), bptt)*loss.cpu().detach()
+
+                    total_positional_losses_recorded += torch.ones(bptt) if single_eval_pos is None else \
+                        nn.functional.one_hot(torch.tensor(single_eval_pos), bptt)
+
+            before_get_batch = time.time()
+        return total_loss / steps_per_epoch, (
+                    total_positional_losses / total_positional_losses_recorded).tolist(), time_to_get_batch, forward_time, step_time
+
+    best_val_loss = float("inf")
+    best_model = None
+    total_loss = float('inf')
+    total_positional_losses = float('inf')
+    try:
+        for epoch in (range(1, epochs + 1) if epochs is not None else itertools.count(1)):
+
+            epoch_start_time = time.time()
+            if train_mixed_precision:
+                with autocast():
+                    total_loss, total_positional_losses, time_to_get_batch, forward_time, step_time = train_step()
+            else:
+                total_loss, total_positional_losses, time_to_get_batch, forward_time, step_time = train_step()
+            if hasattr(dl, 'validate') and epoch % validation_period == 0:
+                with torch.no_grad():
+                    val_score = dl.validate(model)
+            else:
+                val_score = None
+
+            if verbose:
+                print('-' * 89)
+                print(
+                    f'| end of epoch {epoch:3d} | time: {(time.time() - epoch_start_time):5.2f}s | mean loss {total_loss:5.2f} | '
+                    f"pos losses {','.join([f'{l:5.2f}' for l in total_positional_losses])}, lr {scheduler.get_last_lr()[0]}"
+                    f' data time {time_to_get_batch:5.2f} step time {step_time:5.2f}'
+                    f' forward time {forward_time:5.2f}' + (f'val score {val_score}' if val_score is not None else ''))
+                print('-' * 89)
+
+            # stepping with wallclock time based scheduler
+            current_time = time.time()
+            if epoch_callback is not None and rank == 0:
+                epoch_callback(model, epoch / epochs if total_available_time_in_s is None else  # noqa
+                (current_time - start_of_training) / total_available_time_in_s  # noqa
+                               )
+            if epochs is None and (current_time - start_of_training) > total_available_time_in_s:  # noqa
+                break
+            if epochs is None:
+                scheduler.step((current_time - epoch_start_time) / total_available_time_in_s * 100)
+            else:
+                scheduler.step()
+    except KeyboardInterrupt:
+        pass
+
+    return total_loss, total_positional_losses, model.to('cpu'), dl
+
+def _parse_args(config_parser, parser):
+    # Do we have a config file to parse?
+    args_config, remaining = config_parser.parse_known_args()
+    if args_config.config:
+        with open(args_config.config, 'r') as f:
+            cfg = yaml.safe_load(f)
+            parser.set_defaults(**cfg)
+
+    # The main arg parser parses the rest of the args, the usual
+    # defaults will have been overridden if config file specified.
+    args = parser.parse_args(remaining)
+
+    # Cache the args as a text string to save them in the output dir later
+    args_text = yaml.safe_dump(args.__dict__, default_flow_style=False)
+    return args, args_text
+
+
+if __name__ == '__main__':
+    config_parser = argparse.ArgumentParser(description='Only used as a first parser for the config file path.')
+    config_parser.add_argument('--config')
+    parser = argparse.ArgumentParser()
+    parser.add_argument('prior')
+    parser.add_argument('--loss_function', default='barnll')
+    # Optional Arg's for `--loss_function barnll`
+    parser.add_argument('--min_y', type=float, help='barnll can only model y in strict ranges, this is the minimum y can take.')
+    parser.add_argument('--max_y', type=float, help='barnll can only model y in strict ranges, this is the maximum y can take.')
+    parser.add_argument('--num_buckets', default=100, type=int)
+    #parser.add_argument('--num_features', default=None, type=int, help='Specify depending on the prior.')
+    parser.add_argument("--extra_prior_kwargs_dict", default={'fuse_x_y': False}, dest="extra_prior_kwargs_dict", action=StoreDictKeyPair, nargs="+", metavar="KEY=VAL", help='Specify depending on the prior.')
+    parser.add_argument('--encoder', default='linear', type=str, help='Specify depending on the prior.')
+    parser.add_argument('--y_encoder', default='linear', type=str, help='Specify depending on the prior. You should specify this if you do not fuse x and y.')
+    parser.add_argument('--pos_encoder', default='sinus', type=str, help='Specify depending on the prior.')
+    parser.add_argument('--bptt', default=10, type=int)
+    parser.add_argument('--epochs', default=200, type=int)
+    parser.add_argument('--warmup_epochs', default=50, type=int)
+    parser.add_argument('--validation_period', default=10, type=int)
+    parser.add_argument('--permutation_invariant_max_eval_pos', default=None, type=int, help='Set this to an int to ')
+    parser.add_argument('--permutation_invariant_sampling', default='weighted', help="Only relevant if --permutation_invariant_max_eval_pos is set.")
+
+    # these can likely be mostly left at defaults
+    parser.add_argument('--emsize', default=512, type=int) # sometimes even larger is better e.g. 1024
+    parser.add_argument('--nlayers', default=6, type=int)
+    parser.add_argument('--nhid', default=None, type=int) # 2*emsize is the default
+    parser.add_argument('--nhead', default=4, type=int) # nhead = emsize / 64 in the original paper
+    parser.add_argument('--dropout', default=.0, type=float)
+    parser.add_argument('--steps_per_epoch', default=10, type=int)
+    parser.add_argument('--batch_size', default=1000, type=int)
+    parser.add_argument('--lr', '--learning_rate', default=.001, type=float) # try also .0003, .0001, go lower with lower batch size
+
+    args, _ = _parse_args(config_parser, parser)
+
+    if args.nhid is None:
+        args.nhid = 2*args.emsize
+
+    prior = args.__dict__.pop('prior')
+
+    if prior == 'gp':
+        prior = priors.fast_gp.DataLoader
+    elif prior == 'ridge':
+        prior = priors.ridge.DataLoader
+    elif prior == 'stroke':
+        prior = priors.stroke.DataLoader
+    elif prior == 'mix_gp':
+        prior = priors.fast_gp_mix.DataLoader
+    else:
+        raise NotImplementedError(f'Prior == {prior}.')
+
+    loss_function = args.__dict__.pop('loss_function')
+
+    criterion = nn.GaussianNLLLoss(reduction='none', full=True)
+    classificiation_criterion = nn.CrossEntropyLoss(reduction='none')
+    num_buckets = args.__dict__.pop('num_buckets')
+    max_y = args.__dict__.pop('max_y')
+    min_y = args.__dict__.pop('min_y')
+    # criterion = nn.MSELoss(reduction='none')
+
+    def get_y_sample():
+        dl = prior(num_steps=1, batch_size=args.batch_size * args.steps_per_epoch, seq_len=args.bptt, device=device,
+                   **args.extra_prior_kwargs_dict)
+        y_sample = next(iter(dl))[-1]
+        print(f'Creating Bar distribution with borders from y sample of size {y_sample.numel()}')
+        return y_sample
+
+    if loss_function == 'ce':
+        criterion = nn.CrossEntropyLoss(reduction='none')
+    elif loss_function == 'gaussnll':
+        criterion = nn.GaussianNLLLoss(reduction='none', full=True)
+    elif loss_function == 'mse':
+        criterion = nn.MSELoss(reduction='none')
+    elif loss_function == 'barnll':
+        criterion = BarDistribution(borders=get_bucket_limits(num_buckets, full_range=(min_y,max_y)))
+    elif loss_function == 'adaptivebarnll':
+        borders = get_bucket_limits(num_buckets, ys=get_y_sample(), full_range=(min_y,max_y))
+        criterion = BarDistribution(borders=borders)
+    elif loss_function == 'adaptivefullsupportbarnll':
+        assert min_y is None and max_y is None, "Please do not specify `min_y` and `max_y` with `unboundedadaptivebarnll`."
+        borders = get_bucket_limits(num_buckets, ys=get_y_sample())
+        criterion = FullSupportBarDistribution(borders=borders)
+    else:
+        raise NotImplementedError(f'loss_function == {loss_function}.')
+
+
+
+    encoder = args.__dict__.pop('encoder')
+    y_encoder = args.__dict__.pop('y_encoder')
+
+    def get_encoder_generator(encoder):
+        if encoder == 'linear':
+            encoder_generator = encoders.Linear
+        elif encoder == 'mlp':
+            encoder_generator = encoders.MLP
+        elif encoder == 'positional':
+            encoder_generator = encoders.Positional
+        else:
+            raise NotImplementedError(f'A {encoder} encoder is not valid.')
+        return encoder_generator
+
+    encoder_generator = get_encoder_generator(encoder)
+    y_encoder_generator = get_encoder_generator(y_encoder)
+
+    pos_encoder = args.__dict__.pop('pos_encoder')
+
+    if pos_encoder == 'none':
+        pos_encoder_generator = None
+    elif pos_encoder == 'sinus':
+        pos_encoder_generator = positional_encodings.PositionalEncoding
+    elif pos_encoder == 'learned':
+        pos_encoder_generator = positional_encodings.LearnedPositionalEncoding
+    elif pos_encoder == 'paired_scrambled_learned':
+        pos_encoder_generator = positional_encodings.PairedScrambledPositionalEncodings
+    else:
+        raise NotImplementedError(f'pos_encoer == {pos_encoder} is not valid.')
+
+    permutation_invariant_max_eval_pos = args.__dict__.pop('permutation_invariant_max_eval_pos')
+    permutation_invariant_sampling = args.__dict__.pop('permutation_invariant_sampling')
+    if permutation_invariant_max_eval_pos is not None:
+        if permutation_invariant_sampling == 'weighted':
+            get_sampler = get_weighted_single_eval_pos_sampler
+        elif permutation_invariant_sampling == 'uniform':
+            get_sampler = get_uniform_single_eval_pos_sampler
+        else:
+            raise ValueError()
+        args.__dict__['single_eval_pos_gen'] = get_sampler(permutation_invariant_max_eval_pos)
+
+
+    print("ARGS for `train`:", args.__dict__)
+
+    train(prior, criterion, encoder_generator,
+          y_encoder_generator=y_encoder_generator, pos_encoder_generator=pos_encoder_generator,
+          **args.__dict__)
+

TabPFN/transformer.py

@@ -0,0 +1,226 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import Module, TransformerEncoder
+
+from layer import TransformerEncoderLayer, _get_activation_fn
+from utils import SeqBN, bool_mask_to_att_mask
+
+
+
+class TransformerModel(nn.Module):
+    def __init__(self, encoder, n_out, ninp, nhead, nhid, nlayers, dropout=0.0, style_encoder=None, y_encoder=None,
+                 pos_encoder=None, decoder=None, input_normalization=False, init_method=None, pre_norm=False,
+                 activation='gelu', recompute_attn=False, num_global_att_tokens=0, full_attention=False,
+                 all_layers_same_init=True):
+        super().__init__()
+        self.model_type = 'Transformer'
+        encoder_layer_creator = lambda: TransformerEncoderLayer(ninp, nhead, nhid, dropout, activation=activation,
+                                                                pre_norm=pre_norm, recompute_attn=recompute_attn)
+        self.transformer_encoder = TransformerEncoder(encoder_layer_creator(), nlayers)\
+            if all_layers_same_init else TransformerEncoderDiffInit(encoder_layer_creator, nlayers)
+        self.ninp = ninp
+        self.encoder = encoder
+        self.y_encoder = y_encoder
+        self.pos_encoder = pos_encoder
+        self.decoder = decoder(ninp, nhid, n_out) if decoder is not None else nn.Sequential(nn.Linear(ninp, nhid), nn.GELU(), nn.Linear(nhid, n_out))
+        self.input_ln = SeqBN(ninp) if input_normalization else None
+        self.style_encoder = style_encoder
+        self.init_method = init_method
+        if num_global_att_tokens is not None:
+            assert not full_attention
+        self.global_att_embeddings = nn.Embedding(num_global_att_tokens, ninp) if num_global_att_tokens else None
+        self.full_attention = full_attention
+
+        self.n_out = n_out
+        self.nhid = nhid
+
+        self.init_weights()
+
+    @staticmethod
+    def generate_square_subsequent_mask(sz):
+        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+        return bool_mask_to_att_mask(mask)
+
+    @staticmethod
+    def generate_D_q_matrix(sz, query_size):
+        train_size = sz-query_size
+        mask = torch.zeros(sz,sz) == 0
+        mask[:,train_size:].zero_()
+        mask |= torch.eye(sz) == 1
+        return bool_mask_to_att_mask(mask)
+
+    @staticmethod
+    def generate_global_att_query_matrix(num_global_att_tokens, seq_len, num_query_tokens):
+        train_size = seq_len + num_global_att_tokens - num_query_tokens
+        sz = seq_len + num_global_att_tokens
+        mask = torch.zeros(num_query_tokens, sz) == 0
+        mask[:,train_size:].zero_()
+        mask[:,train_size:] |= torch.eye(num_query_tokens) == 1
+        return bool_mask_to_att_mask(mask)
+
+    @staticmethod
+    def generate_global_att_trainset_matrix(num_global_att_tokens, seq_len, num_query_tokens):
+        train_size = seq_len + num_global_att_tokens - num_query_tokens
+        trainset_size = seq_len - num_query_tokens
+        mask = torch.zeros(trainset_size, num_global_att_tokens) == 0
+        #mask[:,num_global_att_tokens:].zero_()
+        #mask[:,num_global_att_tokens:] |= torch.eye(trainset_size) == 1
+        return bool_mask_to_att_mask(mask)
+
+    @staticmethod
+    def generate_global_att_globaltokens_matrix(num_global_att_tokens, seq_len, num_query_tokens):
+        mask = torch.zeros(num_global_att_tokens, num_global_att_tokens+seq_len-num_query_tokens) == 0
+        return bool_mask_to_att_mask(mask)
+
+    def init_weights(self):
+        initrange = 1.
+        # if isinstance(self.encoder,EmbeddingEncoder):
+        #    self.encoder.weight.data.uniform_(-initrange, initrange)
+        # self.decoder.bias.data.zero_()
+        # self.decoder.weight.data.uniform_(-initrange, initrange)
+        if self.init_method is not None:
+            self.apply(self.init_method)
+        for layer in self.transformer_encoder.layers:
+            nn.init.zeros_(layer.linear2.weight)
+            nn.init.zeros_(layer.linear2.bias)
+            attns = layer.self_attn if isinstance(layer.self_attn, nn.ModuleList) else [layer.self_attn]
+            for attn in attns:
+                nn.init.zeros_(attn.out_proj.weight)
+                nn.init.zeros_(attn.out_proj.bias)
+
+    def forward(self, src, src_mask=None, single_eval_pos=None):
+        assert isinstance(src, tuple), 'fuse_x_y is forbidden, that is inputs have to be given as (x,y) or (style,x,y)'
+
+        if len(src) == 2:
+            src = (None,) + src
+
+        style_src, style_src_size = (src[0], (0 if (src[0] is None) else 1))
+        if src_mask is not None: assert self.global_att_embeddings is None or isinstance(src_mask, tuple)
+        if src_mask is None:
+            x_src = src[1]
+            if self.global_att_embeddings is None:
+                full_len = len(x_src) + style_src_size
+                if self.full_attention:
+                    src_mask = bool_mask_to_att_mask(torch.ones((full_len, full_len), dtype=torch.bool)).to(x_src.device)
+                else:
+                    src_mask = self.generate_D_q_matrix(len(x_src) + style_src_size, len(x_src) + style_src_size -single_eval_pos).to(x_src.device)
+            else:
+                src_mask_args = (self.global_att_embeddings.num_embeddings,
+                                 len(x_src) + style_src_size,
+                                 len(x_src) + style_src_size - single_eval_pos)
+                src_mask = (self.generate_global_att_globaltokens_matrix(*src_mask_args).to(x_src.device),
+                            self.generate_global_att_trainset_matrix(*src_mask_args).to(x_src.device),
+                            self.generate_global_att_query_matrix(*src_mask_args).to(x_src.device))
+
+        style_src, x_src, y_src = src
+        x_src = self.encoder(x_src)
+        y_src = self.y_encoder(y_src.unsqueeze(-1) if len(y_src.shape) < len(x_src.shape) else y_src)
+        style_src = self.style_encoder(style_src).unsqueeze(0) if self.style_encoder else torch.tensor([], device=x_src.device)
+        global_src = torch.tensor([], device=x_src.device) if self.global_att_embeddings is None else \
+            self.global_att_embeddings.weight.unsqueeze(1).repeat(1, x_src.shape[1], 1)
+        train_x = x_src[:single_eval_pos] + y_src[:single_eval_pos]
+        src = torch.cat([global_src, style_src, train_x, x_src[single_eval_pos:]], 0)
+
+        if self.input_ln is not None:
+            src = self.input_ln(src)
+
+        if self.pos_encoder is not None:
+            src = self.pos_encoder(src)
+
+        # If we have style input, drop its output
+        output = self.transformer_encoder(src, src_mask)[style_src_size:]
+        output = self.decoder(output)
+        return output[single_eval_pos+(self.global_att_embeddings.num_embeddings if self.global_att_embeddings else 0):]
+
+    @torch.no_grad()
+    def init_from_small_model(self, small_model):
+        assert isinstance(self.decoder, nn.Linear) and isinstance(self.encoder, (nn.Linear, nn.Sequential)) \
+               and isinstance(self.y_encoder, (nn.Linear, nn.Sequential))
+
+        def set_encoder_weights(my_encoder, small_model_encoder):
+            my_encoder_linear, small_encoder_linear = (my_encoder, small_model_encoder) \
+                if isinstance(my_encoder, nn.Linear) else (my_encoder[-1], small_model_encoder[-1])
+            small_in_dim = small_encoder_linear.out_features
+            my_encoder_linear.weight.zero_()
+            my_encoder_linear.bias.zero_()
+            my_encoder_linear.weight[:small_in_dim] = small_encoder_linear.weight
+            my_encoder_linear.bias[:small_in_dim] = small_encoder_linear.bias
+
+        set_encoder_weights(self.encoder, small_model.encoder)
+        set_encoder_weights(self.y_encoder, small_model.y_encoder)
+
+        small_in_dim = small_model.decoder.in_features
+
+        self.decoder.weight[:, :small_in_dim] = small_model.decoder.weight
+        self.decoder.bias = small_model.decoder.bias
+
+        for my_layer, small_layer in zip(self.transformer_encoder.layers, small_model.transformer_encoder.layers):
+            small_hid_dim = small_layer.linear1.out_features
+            my_in_dim = my_layer.linear1.in_features
+
+            # packed along q,k,v order in first dim
+            my_in_proj_w = my_layer.self_attn.in_proj_weight
+            small_in_proj_w = small_layer.self_attn.in_proj_weight
+
+            my_in_proj_w.view(3, my_in_dim, my_in_dim)[:, :small_in_dim, :small_in_dim] = small_in_proj_w.view(3,
+                                                                                                               small_in_dim,
+                                                                                                               small_in_dim)
+            my_layer.self_attn.in_proj_bias.view(3, my_in_dim)[:,
+            :small_in_dim] = small_layer.self_attn.in_proj_bias.view(3, small_in_dim)
+
+            my_layer.self_attn.out_proj.weight[:small_in_dim, :small_in_dim] = small_layer.self_attn.out_proj.weight
+            my_layer.self_attn.out_proj.bias[:small_in_dim] = small_layer.self_attn.out_proj.bias
+
+            my_layer.linear1.weight[:small_hid_dim, :small_in_dim] = small_layer.linear1.weight
+            my_layer.linear1.bias[:small_hid_dim] = small_layer.linear1.bias
+
+            my_layer.linear2.weight[:small_in_dim, :small_hid_dim] = small_layer.linear2.weight
+            my_layer.linear2.bias[:small_in_dim] = small_layer.linear2.bias
+
+            my_layer.norm1.weight[:small_in_dim] = math.sqrt(small_in_dim / my_in_dim) * small_layer.norm1.weight
+            my_layer.norm2.weight[:small_in_dim] = math.sqrt(small_in_dim / my_in_dim) * small_layer.norm2.weight
+
+            my_layer.norm1.bias[:small_in_dim] = small_layer.norm1.bias
+            my_layer.norm2.bias[:small_in_dim] = small_layer.norm2.bias
+
+
+class TransformerEncoderDiffInit(Module):
+    r"""TransformerEncoder is a stack of N encoder layers
+
+    Args:
+        encoder_layer_creator: a function generating objects of TransformerEncoderLayer class without args (required).
+        num_layers: the number of sub-encoder-layers in the encoder (required).
+        norm: the layer normalization component (optional).
+    """
+    __constants__ = ['norm']
+
+    def __init__(self, encoder_layer_creator, num_layers, norm=None):
+        super().__init__()
+        self.layers = nn.ModuleList([encoder_layer_creator() for _ in range(num_layers)])
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
+        r"""Pass the input through the encoder layers in turn.
+
+        Args:
+            src: the sequence to the encoder (required).
+            mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        output = src
+
+        for mod in self.layers:
+            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output

TabPFN/utils.py

@@ -0,0 +1,236 @@
+import os
+import math
+import argparse
+import random
+import datetime
+
+import torch
+from torch import nn
+from torch.optim.lr_scheduler import LambdaLR
+import numpy as np
+
+# copied from huggingface
+def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, num_cycles=0.5, last_epoch=-1):
+    """ Create a schedule with a learning rate that decreases following the
+    values of the cosine function between 0 and `pi * cycles` after a warmup
+    period during which it increases linearly between 0 and 1.
+    """
+
+    def lr_lambda(current_step):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))
+
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+# copied from huggingface
+def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
+    """
+    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after
+    a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
+
+    Args:
+        optimizer (:class:`~torch.optim.Optimizer`):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (:obj:`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (:obj:`int`):
+            The total number of training steps.
+        last_epoch (:obj:`int`, `optional`, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+
+    def lr_lambda(current_step: int):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        return max(
+            0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
+        )
+
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def get_openai_lr(transformer_model):
+    num_params = sum(p.numel() for p in transformer_model.parameters())
+    return 0.003239 - 0.0001395 * math.log(num_params)
+
+
+def get_weighted_single_eval_pos_sampler(max_len):
+    """
+    This gives a sampler that can be used for `single_eval_pos` which yields good performance for all positions p,
+    where p <= `max_len`. At most `max_len` - 1 examples are shown to the Transformer.
+    :return: Sampler that can be fed to `train()` as `single_eval_pos_gen`.
+    """
+    return lambda: random.choices(range(max_len), [1 / (max_len - i) for i in range(max_len)])[0]
+
+
+def get_uniform_single_eval_pos_sampler(max_len, min_len=0):
+    """
+    Just sample any evaluation position with the same weight
+    :return: Sampler that can be fed to `train()` as `single_eval_pos_gen`.
+    """
+    return lambda: random.choices(range(min_len, max_len))[0]
+
+
+class SeqBN(nn.Module):
+    def __init__(self, d_model):
+        super().__init__()
+        self.bn = nn.BatchNorm1d(d_model)
+        self.d_model = d_model
+
+    def forward(self, x):
+        assert self.d_model == x.shape[-1]
+        flat_x = x.view(-1, self.d_model)
+        flat_x = self.bn(flat_x)
+        return flat_x.view(*x.shape)
+
+
+def set_locals_in_self(locals):
+    self = locals['self']
+    for var_name, val in locals.items():
+        if var_name != 'self': setattr(self, var_name, val)
+
+
+default_device = 'cuda:0' if torch.cuda.is_available() else 'cpu:0'
+
+
+# Copied from StackOverflow, but we do an eval on the values additionally
+class StoreDictKeyPair(argparse.Action):
+    def __init__(self, option_strings, dest, nargs=None, **kwargs):
+        self._nargs = nargs
+        super(StoreDictKeyPair, self).__init__(option_strings, dest, nargs=nargs, **kwargs)
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        my_dict = {}
+        for kv in values:
+            k, v = kv.split("=")
+            try:
+                my_dict[k] = eval(v)
+            except NameError:
+                my_dict[k] = v
+        setattr(namespace, self.dest, my_dict)
+        print("dict values: {}".format(my_dict))
+
+def get_nan_value(v, set_value_to_nan=0.0):
+    if random.random() < set_value_to_nan:
+        return v
+    else:
+        return random.choice([-999, 0, 1, 999])
+
+def to_ranking(data):
+    x = (data >= data.unsqueeze(-3))
+    x = x.sum(0)
+    return x
+# TODO: Is there a better way to do this?
+#   1. Cmparing to unique elements: When all values are different we still get quadratic blowup
+#   2. Argsort(Argsort()) returns ranking, but with duplicate values there is an ordering which is problematic
+#   3. Argsort(Argsort(Unique))->Scatter seems a bit complicated, doesn't have quadratic blowup, but how fast?
+def to_ranking_low_mem(data):
+    x = torch.zeros_like(data)
+    for col in range(data.shape[-1]):
+        x_ = (data[:, :, col] >= data[:, :, col].unsqueeze(-2))
+        x_ = x_.sum(0)
+        x[:, :, col] = x_
+    return x
+
+def nan_handling_missing_for_unknown_reason_value(set_value_to_nan=0.0):
+    return get_nan_value(float('nan'), set_value_to_nan)
+
+def nan_handling_missing_for_no_reason_value(set_value_to_nan=0.0):
+    return get_nan_value(float('-inf'), set_value_to_nan)
+
+def nan_handling_missing_for_a_reason_value(set_value_to_nan=0.0):
+    return get_nan_value(float('inf'), set_value_to_nan)
+
+def torch_nanmean(x, axis=0):
+    num = torch.where(torch.isnan(x), torch.full_like(x, 0), torch.full_like(x, 1)).sum(axis=axis)
+    value = torch.where(torch.isnan(x), torch.full_like(x, 0), x).sum(axis=axis)
+    return value / num
+
+def torch_nanstd(x, axis=0):
+    num = torch.where(torch.isnan(x), torch.full_like(x, 0), torch.full_like(x, 1)).sum(axis=axis)
+    value = torch.where(torch.isnan(x), torch.full_like(x, 0), x).sum(axis=axis)
+    mean = value / num
+    mean_broadcast = torch.repeat_interleave(mean.unsqueeze(axis), x.shape[axis], dim=axis)
+    return torch.sqrt(torch.nansum(torch.square(mean_broadcast - x), axis=axis) / (num - 1))
+
+def normalize_data(data, normalize_positions=-1):
+    if normalize_positions > 0:
+        mean = torch_nanmean(data[:normalize_positions], axis=0)
+        std = torch_nanstd(data[:normalize_positions], axis=0) + .000001
+    else:
+        mean = torch_nanmean(data, axis=0)
+        std = torch_nanstd(data, axis=0) + .000001
+    data = (data - mean) / std
+    data = torch.clip(data, min=-100, max=100)
+
+    return data
+
+def remove_outliers(X, n_sigma=4):
+    # Expects T, B, H
+    assert len(X.shape) == 3, "X must be T,B,H"
+    #for b in range(X.shape[1]):
+        #for col in range(X.shape[2]):
+    data = X
+    data_mean, data_std = torch_nanmean(data, axis=0), torch_nanstd(data, axis=0)
+    cut_off = data_std * n_sigma
+    lower, upper = data_mean - cut_off, data_mean + cut_off
+
+    data_clean = X[:].clone()
+    data_clean[torch.logical_or(data > upper, data < lower)] = np.nan
+    data_mean, data_std = torch_nanmean(data_clean, axis=0), torch_nanstd(data_clean, axis=0)
+    cut_off = data_std * n_sigma
+    lower, upper = data_mean - cut_off, data_mean + cut_off
+
+    X = torch.maximum(-torch.log(1+torch.abs(X)) + lower, X)
+    X = torch.minimum(torch.log(1+torch.abs(X)) + upper, X)
+            # print(ds[1][data < lower, col], ds[1][data > upper, col], ds[1][~np.isnan(data), col].shape, data_mean, data_std)
+    return X
+
+def bool_mask_to_att_mask(mask):
+    return mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+
+def print_on_master_only(is_master):
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+def init_dist(device):
+    if 'SLURM_PROCID' in os.environ and torch.cuda.device_count() > 1:
+        assert device != 'cpu:0'
+        rank = int(os.environ['SLURM_PROCID'])
+        os.environ['MASTER_ADDR'] = 'localhost'
+        os.environ['MASTER_PORT'] = '12355'
+        torch.cuda.set_device(rank)
+        os.environ['CUDA_VISIBLE_DEVICES'] = str(rank)
+        torch.distributed.init_process_group(backend="nccl", init_method="env://", timeout=datetime.timedelta(seconds=20),
+                                             world_size=torch.cuda.device_count(), rank=rank)
+        torch.distributed.barrier()
+        print_on_master_only(rank == 0)
+        print(f"Distributed training on {torch.cuda.device_count()} GPUs, this is rank {rank}, "
+              "only I can print, but when using print(..., force=True) it will print on all ranks.")
+
+        return True, rank, f'cuda:{rank}'
+    else:
+        print('Not using distributed')
+        # will not change any of the behavior of print, but allows putting the force=True in the print calls
+        print_on_master_only(True)
+        return False, 0, device
+
+# NOP function for python with statements (x = NOP(); with x:)
+class NOP():
+    def __enter__(self):
+        pass
+    def __exit__(self, type, value, traceback):
+        pass

app.py

@@ -0,0 +1,96 @@
+import sys
+tabpfn_path = 'TabPFN'
+sys.path.insert(0, tabpfn_path) # our submodule of the TabPFN repo (at 045c8400203ebd062346970b4f2c0ccda5a40618)
+from TabPFN.scripts.transformer_prediction_interface import TabPFNClassifier
+
+import numpy as np
+import pandas as pd
+import torch
+import gradio as gr
+import openml
+
+
+def compute(table: np.array):
+    vfunc = np.vectorize(lambda s: len(s))
+    non_empty_row_mask = (vfunc(table).sum(1) != 0)
+    table = table[non_empty_row_mask]
+    empty_mask = table == ''
+    empty_inds = np.where(empty_mask)
+    if not len(empty_inds[0]):
+        return "**Please leave at least one field blank for prediction.**", None
+    if not np.all(empty_inds[1][0] == empty_inds[1]):
+        return "**Please only leave fields of one column blank for prediction.**", None
+    y_column = empty_inds[1][0]
+    eval_lines = empty_inds[0]
+
+    train_table = np.delete(table, eval_lines, axis=0)
+    eval_table = table[eval_lines]
+
+    try:
+        x_train = torch.tensor(np.delete(train_table, y_column, axis=1).astype(np.float32))
+        x_eval = torch.tensor(np.delete(eval_table, y_column, axis=1).astype(np.float32))
+
+        y_train = train_table[:, y_column]
+    except ValueError:
+        return "**Please only add numbers (to the inputs) or leave fields empty.**", None
+
+    classifier = TabPFNClassifier(base_path=tabpfn_path, device='cpu')
+    classifier.fit(x_train, y_train)
+    y_eval, p_eval = classifier.predict(x_eval, return_winning_probability=True)
+
+    # print(file, type(file))
+    out_table = table.copy().astype(str)
+    out_table[eval_lines, y_column] = [f"{y_e} (p={p_e:.2f})" for y_e, p_e in zip(y_eval, p_eval)]
+    return None, out_table
+
+
+def upload_file(file):
+    if file.name.endswith('.arff'):
+        dataset = openml.datasets.OpenMLDataset('t', 'test', data_file=file.name)
+        X_, _, categorical_indicator_, attribute_names_ = dataset.get_data(
+            dataset_format="array"
+        )
+        df = pd.DataFrame(X_, columns=attribute_names_)
+        return df
+    elif file.name.endswith('.csv') or file.name.endswith('.data'):
+        df = pd.read_csv(file.name, header=None)
+        df.columns = np.arange(len(df.columns))
+        print(df)
+        return df
+
+
+example = \
+    [
+        [1, 2, 1],
+        [2, 1, 1],
+        [1, 1, 1],
+        [2, 2, 2],
+        [3, 4, 2],
+        [3, 2, 2],
+        [2, 3, '']
+    ]
+
+with gr.Blocks() as demo:
+    gr.Markdown("""This demo allows you to play with the **TabPFN**.
+    You can either change the table manually (we have filled it with a toy benchmark, sum up to 3 has label 1 and over that label 2).
+    The network predicts fields you leave empty. Only one column can have empty entries that are predicted.
+    Please, provide everything but the label column as numeric values. It is ok to encode classes as integers.
+    """)
+    inp_table = gr.DataFrame(type='numpy', value=example, headers=[''] * 3)
+    inp_file = gr.File(
+        label='Drop either a .csv (without header, only numeric values for all but the labels) or a .arff file.')
+    examples = gr.Examples(examples=['iris.csv', 'balance-scale.arff'],
+                           inputs=[inp_file],
+                           outputs=[inp_table],
+                           fn=upload_file,
+                           cache_examples=True)
+    btn = gr.Button("Predict Empty Table Cells")
+
+    inp_file.change(fn=upload_file, inputs=inp_file, outputs=inp_table)
+
+    out_text = gr.Markdown()
+    out_table = gr.DataFrame()
+
+    btn.click(fn=compute, inputs=inp_table, outputs=[out_text, out_table])
+
+demo.launch()

balance-scale.arff

@@ -0,0 +1,694 @@
+%1. Title: Balance Scale Weight & Distance Database
+%
+%2. Source Information: 
+%    (a) Source: Generated to model psychological experiments reported
+%		by Siegler, R. S. (1976).  Three Aspects of Cognitive
+%		Development.  Cognitive Psychology, 8, 481-520.
+%    (b) Donor: Tim Hume (hume@ics.uci.edu)
+%    (c) Date: 22 April 1994
+%
+%3. Past Usage: (possibly different formats of this data)
+%   - Publications
+%	1. Klahr, D., & Siegler, R.S. (1978).  The Representation of
+%	   Children's Knowledge.  In H. W. Reese & L. P. Lipsitt (Eds.),
+%	   Advances in Child Development and Behavior, pp. 61-116.  New
+%	   York: Academic Press 
+%	2. Langley,P. (1987).  A General Theory of Discrimination
+%	   Learning.  In D. Klahr, P. Langley, & R. Neches (Eds.),
+%	   Production System Models of Learning and Development, pp.
+%	   99-161. Cambridge, MA: MIT Press
+%	3. Newell, A. (1990).  Unified Theories of Cognition.
+%	   Cambridge, MA: Harvard University Press
+%	4. McClelland, J.L. (1988).  Parallel Distibuted Processing:
+%	   Implications for Cognition and Development.  Technical
+%	   Report AIP-47, Department of Psychology, Carnegie-Mellon
+%	   University 
+%	5. Shultz, T., Mareschal, D., & Schmidt, W. (1994).  Modeling
+%	   Cognitive Development on Balance Scale Phenomena. Machine
+%	   Learning, Vol. 16, pp. 59-88.
+%
+%4. Relevant Information: 
+%	This data set was generated to model psychological
+%	experimental results.  Each example is classified as having the
+%	balance scale tip to the right, tip to the left, or be
+%	balanced.  The attributes are the left weight, the left
+%	distance, the right weight, and the right distance.  The
+%	correct way to find the class is the greater of 
+%	(left-distance * left-weight) and (right-distance *
+%	right-weight).  If they are equal, it is balanced.
+%
+%5. Number of Instances: 625 (49 balanced, 288 left, 288 right)
+%
+%6. Number of Attributes: 4 (numeric) + class name = 5
+%
+%7. Attribute Information:
+%	1. Class Name: 3 (L, B, R)
+%	2. Left-Weight: 5 (1, 2, 3, 4, 5)
+%	3. Left-Distance: 5 (1, 2, 3, 4, 5)
+%	4. Right-Weight: 5 (1, 2, 3, 4, 5)
+%	5. Right-Distance: 5 (1, 2, 3, 4, 5)
+%
+%8. Missing Attribute Values: 
+%	none
+%
+%9. Class Distribution: 
+%   1. 46.08 percent are L
+%   2. 07.84 percent are B
+%   3. 46.08 percent are R
+%
+
+@relation balance-scale
+@attribute 'left-weight' real
+@attribute 'left-distance' real
+@attribute 'right-weight' real
+@attribute 'right-distance' real
+@attribute 'class' { L, B, R}
+@data
+1,1,1,1,B
+1,1,1,2,R
+1,1,1,3,R
+1,1,1,4,R
+1,1,1,5,R
+1,1,2,1,R
+1,1,2,2,R
+1,1,2,3,R
+1,1,2,4,R
+1,1,2,5,R
+1,1,3,1,R
+1,1,3,2,R
+1,1,3,3,R
+1,1,3,4,R
+1,1,3,5,R
+1,1,4,1,R
+1,1,4,2,R
+1,1,4,3,R
+1,1,4,4,R
+1,1,4,5,R
+1,1,5,1,R
+1,1,5,2,R
+1,1,5,3,R
+1,1,5,4,R
+1,1,5,5,R
+1,2,1,1,L
+1,2,1,2,B
+1,2,1,3,R
+1,2,1,4,R
+1,2,1,5,R
+1,2,2,1,B
+1,2,2,2,R
+1,2,2,3,R
+1,2,2,4,R
+1,2,2,5,R
+1,2,3,1,R
+1,2,3,2,R
+1,2,3,3,R
+1,2,3,4,R
+1,2,3,5,R
+1,2,4,1,R
+1,2,4,2,R
+1,2,4,3,R
+1,2,4,4,R
+1,2,4,5,R
+1,2,5,1,R
+1,2,5,2,R
+1,2,5,3,R
+1,2,5,4,R
+1,2,5,5,R
+1,3,1,1,L
+1,3,1,2,L
+1,3,1,3,B
+1,3,1,4,R
+1,3,1,5,R
+1,3,2,1,L
+1,3,2,2,R
+1,3,2,3,R
+1,3,2,4,R
+1,3,2,5,R
+1,3,3,1,B
+1,3,3,2,R
+1,3,3,3,R
+1,3,3,4,R
+1,3,3,5,R
+1,3,4,1,R
+1,3,4,2,R
+1,3,4,3,R
+1,3,4,4,R
+1,3,4,5,R
+1,3,5,1,R
+1,3,5,2,R
+1,3,5,3,R
+1,3,5,4,R
+1,3,5,5,R
+1,4,1,1,L
+1,4,1,2,L
+1,4,1,3,L
+1,4,1,4,B
+1,4,1,5,R
+1,4,2,1,L
+1,4,2,2,B
+1,4,2,3,R
+1,4,2,4,R
+1,4,2,5,R
+1,4,3,1,L
+1,4,3,2,R
+1,4,3,3,R
+1,4,3,4,R
+1,4,3,5,R
+1,4,4,1,B
+1,4,4,2,R
+1,4,4,3,R
+1,4,4,4,R
+1,4,4,5,R
+1,4,5,1,R
+1,4,5,2,R
+1,4,5,3,R
+1,4,5,4,R
+1,4,5,5,R
+1,5,1,1,L
+1,5,1,2,L
+1,5,1,3,L
+1,5,1,4,L
+1,5,1,5,B
+1,5,2,1,L
+1,5,2,2,L
+1,5,2,3,R
+1,5,2,4,R
+1,5,2,5,R
+1,5,3,1,L
+1,5,3,2,R
+1,5,3,3,R
+1,5,3,4,R
+1,5,3,5,R
+1,5,4,1,L
+1,5,4,2,R
+1,5,4,3,R
+1,5,4,4,R
+1,5,4,5,R
+1,5,5,1,B
+1,5,5,2,R
+1,5,5,3,R
+1,5,5,4,R
+1,5,5,5,R
+2,1,1,1,L
+2,1,1,2,B
+2,1,1,3,R
+2,1,1,4,R
+2,1,1,5,R
+2,1,2,1,B
+2,1,2,2,R
+2,1,2,3,R
+2,1,2,4,R
+2,1,2,5,R
+2,1,3,1,R
+2,1,3,2,R
+2,1,3,3,R
+2,1,3,4,R
+2,1,3,5,R
+2,1,4,1,R
+2,1,4,2,R
+2,1,4,3,R
+2,1,4,4,R
+2,1,4,5,R
+2,1,5,1,R
+2,1,5,2,R
+2,1,5,3,R
+2,1,5,4,R
+2,1,5,5,R
+2,2,1,1,L
+2,2,1,2,L
+2,2,1,3,L
+2,2,1,4,B
+2,2,1,5,R
+2,2,2,1,L
+2,2,2,2,B
+2,2,2,3,R
+2,2,2,4,R
+2,2,2,5,R
+2,2,3,1,L
+2,2,3,2,R
+2,2,3,3,R
+2,2,3,4,R
+2,2,3,5,R
+2,2,4,1,B
+2,2,4,2,R
+2,2,4,3,R
+2,2,4,4,R
+2,2,4,5,R
+2,2,5,1,R
+2,2,5,2,R
+2,2,5,3,R
+2,2,5,4,R
+2,2,5,5,R
+2,3,1,1,L
+2,3,1,2,L
+2,3,1,3,L
+2,3,1,4,L
+2,3,1,5,L
+2,3,2,1,L
+2,3,2,2,L
+2,3,2,3,B
+2,3,2,4,R
+2,3,2,5,R
+2,3,3,1,L
+2,3,3,2,B
+2,3,3,3,R
+2,3,3,4,R
+2,3,3,5,R
+2,3,4,1,L
+2,3,4,2,R
+2,3,4,3,R
+2,3,4,4,R
+2,3,4,5,R
+2,3,5,1,L
+2,3,5,2,R
+2,3,5,3,R
+2,3,5,4,R
+2,3,5,5,R
+2,4,1,1,L
+2,4,1,2,L
+2,4,1,3,L
+2,4,1,4,L
+2,4,1,5,L
+2,4,2,1,L
+2,4,2,2,L
+2,4,2,3,L
+2,4,2,4,B
+2,4,2,5,R
+2,4,3,1,L
+2,4,3,2,L
+2,4,3,3,R
+2,4,3,4,R
+2,4,3,5,R
+2,4,4,1,L
+2,4,4,2,B
+2,4,4,3,R
+2,4,4,4,R
+2,4,4,5,R
+2,4,5,1,L
+2,4,5,2,R
+2,4,5,3,R
+2,4,5,4,R
+2,4,5,5,R
+2,5,1,1,L
+2,5,1,2,L
+2,5,1,3,L
+2,5,1,4,L
+2,5,1,5,L
+2,5,2,1,L
+2,5,2,2,L
+2,5,2,3,L
+2,5,2,4,L
+2,5,2,5,B
+2,5,3,1,L
+2,5,3,2,L
+2,5,3,3,L
+2,5,3,4,R
+2,5,3,5,R
+2,5,4,1,L
+2,5,4,2,L
+2,5,4,3,R
+2,5,4,4,R
+2,5,4,5,R
+2,5,5,1,L
+2,5,5,2,B
+2,5,5,3,R
+2,5,5,4,R
+2,5,5,5,R
+3,1,1,1,L
+3,1,1,2,L
+3,1,1,3,B
+3,1,1,4,R
+3,1,1,5,R
+3,1,2,1,L
+3,1,2,2,R
+3,1,2,3,R
+3,1,2,4,R
+3,1,2,5,R
+3,1,3,1,B
+3,1,3,2,R
+3,1,3,3,R
+3,1,3,4,R
+3,1,3,5,R
+3,1,4,1,R
+3,1,4,2,R
+3,1,4,3,R
+3,1,4,4,R
+3,1,4,5,R
+3,1,5,1,R
+3,1,5,2,R
+3,1,5,3,R
+3,1,5,4,R
+3,1,5,5,R
+3,2,1,1,L
+3,2,1,2,L
+3,2,1,3,L
+3,2,1,4,L
+3,2,1,5,L
+3,2,2,1,L
+3,2,2,2,L
+3,2,2,3,B
+3,2,2,4,R
+3,2,2,5,R
+3,2,3,1,L
+3,2,3,2,B
+3,2,3,3,R
+3,2,3,4,R
+3,2,3,5,R
+3,2,4,1,L
+3,2,4,2,R
+3,2,4,3,R
+3,2,4,4,R
+3,2,4,5,R
+3,2,5,1,L
+3,2,5,2,R
+3,2,5,3,R
+3,2,5,4,R
+3,2,5,5,R
+3,3,1,1,L
+3,3,1,2,L
+3,3,1,3,L
+3,3,1,4,L
+3,3,1,5,L
+3,3,2,1,L
+3,3,2,2,L
+3,3,2,3,L
+3,3,2,4,L
+3,3,2,5,R
+3,3,3,1,L
+3,3,3,2,L
+3,3,3,3,B
+3,3,3,4,R
+3,3,3,5,R
+3,3,4,1,L
+3,3,4,2,L
+3,3,4,3,R
+3,3,4,4,R
+3,3,4,5,R
+3,3,5,1,L
+3,3,5,2,R
+3,3,5,3,R
+3,3,5,4,R
+3,3,5,5,R
+3,4,1,1,L
+3,4,1,2,L
+3,4,1,3,L
+3,4,1,4,L
+3,4,1,5,L
+3,4,2,1,L
+3,4,2,2,L
+3,4,2,3,L
+3,4,2,4,L
+3,4,2,5,L
+3,4,3,1,L
+3,4,3,2,L
+3,4,3,3,L
+3,4,3,4,B
+3,4,3,5,R
+3,4,4,1,L
+3,4,4,2,L
+3,4,4,3,B
+3,4,4,4,R
+3,4,4,5,R
+3,4,5,1,L
+3,4,5,2,L
+3,4,5,3,R
+3,4,5,4,R
+3,4,5,5,R
+3,5,1,1,L
+3,5,1,2,L
+3,5,1,3,L
+3,5,1,4,L
+3,5,1,5,L
+3,5,2,1,L
+3,5,2,2,L
+3,5,2,3,L
+3,5,2,4,L
+3,5,2,5,L
+3,5,3,1,L
+3,5,3,2,L
+3,5,3,3,L
+3,5,3,4,L
+3,5,3,5,B
+3,5,4,1,L
+3,5,4,2,L
+3,5,4,3,L
+3,5,4,4,R
+3,5,4,5,R
+3,5,5,1,L
+3,5,5,2,L
+3,5,5,3,B
+3,5,5,4,R
+3,5,5,5,R
+4,1,1,1,L
+4,1,1,2,L
+4,1,1,3,L
+4,1,1,4,B
+4,1,1,5,R
+4,1,2,1,L
+4,1,2,2,B
+4,1,2,3,R
+4,1,2,4,R
+4,1,2,5,R
+4,1,3,1,L
+4,1,3,2,R
+4,1,3,3,R
+4,1,3,4,R
+4,1,3,5,R
+4,1,4,1,B
+4,1,4,2,R
+4,1,4,3,R
+4,1,4,4,R
+4,1,4,5,R
+4,1,5,1,R
+4,1,5,2,R
+4,1,5,3,R
+4,1,5,4,R
+4,1,5,5,R
+4,2,1,1,L
+4,2,1,2,L
+4,2,1,3,L
+4,2,1,4,L
+4,2,1,5,L
+4,2,2,1,L
+4,2,2,2,L
+4,2,2,3,L
+4,2,2,4,B
+4,2,2,5,R
+4,2,3,1,L
+4,2,3,2,L
+4,2,3,3,R
+4,2,3,4,R
+4,2,3,5,R
+4,2,4,1,L
+4,2,4,2,B
+4,2,4,3,R
+4,2,4,4,R
+4,2,4,5,R
+4,2,5,1,L
+4,2,5,2,R
+4,2,5,3,R
+4,2,5,4,R
+4,2,5,5,R
+4,3,1,1,L
+4,3,1,2,L
+4,3,1,3,L
+4,3,1,4,L
+4,3,1,5,L
+4,3,2,1,L
+4,3,2,2,L
+4,3,2,3,L
+4,3,2,4,L
+4,3,2,5,L
+4,3,3,1,L
+4,3,3,2,L
+4,3,3,3,L
+4,3,3,4,B
+4,3,3,5,R
+4,3,4,1,L
+4,3,4,2,L
+4,3,4,3,B
+4,3,4,4,R
+4,3,4,5,R
+4,3,5,1,L
+4,3,5,2,L
+4,3,5,3,R
+4,3,5,4,R
+4,3,5,5,R
+4,4,1,1,L
+4,4,1,2,L
+4,4,1,3,L
+4,4,1,4,L
+4,4,1,5,L
+4,4,2,1,L
+4,4,2,2,L
+4,4,2,3,L
+4,4,2,4,L
+4,4,2,5,L
+4,4,3,1,L
+4,4,3,2,L
+4,4,3,3,L
+4,4,3,4,L
+4,4,3,5,L
+4,4,4,1,L
+4,4,4,2,L
+4,4,4,3,L
+4,4,4,4,B
+4,4,4,5,R
+4,4,5,1,L
+4,4,5,2,L
+4,4,5,3,L
+4,4,5,4,R
+4,4,5,5,R
+4,5,1,1,L
+4,5,1,2,L
+4,5,1,3,L
+4,5,1,4,L
+4,5,1,5,L
+4,5,2,1,L
+4,5,2,2,L
+4,5,2,3,L
+4,5,2,4,L
+4,5,2,5,L
+4,5,3,1,L
+4,5,3,2,L
+4,5,3,3,L
+4,5,3,4,L
+4,5,3,5,L
+4,5,4,1,L
+4,5,4,2,L
+4,5,4,3,L
+4,5,4,4,L
+4,5,4,5,B
+4,5,5,1,L
+4,5,5,2,L
+4,5,5,3,L
+4,5,5,4,B
+4,5,5,5,R
+5,1,1,1,L
+5,1,1,2,L
+5,1,1,3,L
+5,1,1,4,L
+5,1,1,5,B
+5,1,2,1,L
+5,1,2,2,L
+5,1,2,3,R
+5,1,2,4,R
+5,1,2,5,R
+5,1,3,1,L
+5,1,3,2,R
+5,1,3,3,R
+5,1,3,4,R
+5,1,3,5,R
+5,1,4,1,L
+5,1,4,2,R
+5,1,4,3,R
+5,1,4,4,R
+5,1,4,5,R
+5,1,5,1,B
+5,1,5,2,R
+5,1,5,3,R
+5,1,5,4,R
+5,1,5,5,R
+5,2,1,1,L
+5,2,1,2,L
+5,2,1,3,L
+5,2,1,4,L
+5,2,1,5,L
+5,2,2,1,L
+5,2,2,2,L
+5,2,2,3,L
+5,2,2,4,L
+5,2,2,5,B
+5,2,3,1,L
+5,2,3,2,L
+5,2,3,3,L
+5,2,3,4,R
+5,2,3,5,R
+5,2,4,1,L
+5,2,4,2,L
+5,2,4,3,R
+5,2,4,4,R
+5,2,4,5,R
+5,2,5,1,L
+5,2,5,2,B
+5,2,5,3,R
+5,2,5,4,R
+5,2,5,5,R
+5,3,1,1,L
+5,3,1,2,L
+5,3,1,3,L
+5,3,1,4,L
+5,3,1,5,L
+5,3,2,1,L
+5,3,2,2,L
+5,3,2,3,L
+5,3,2,4,L
+5,3,2,5,L
+5,3,3,1,L
+5,3,3,2,L
+5,3,3,3,L
+5,3,3,4,L
+5,3,3,5,B
+5,3,4,1,L
+5,3,4,2,L
+5,3,4,3,L
+5,3,4,4,R
+5,3,4,5,R
+5,3,5,1,L
+5,3,5,2,L
+5,3,5,3,B
+5,3,5,4,R
+5,3,5,5,R
+5,4,1,1,L
+5,4,1,2,L
+5,4,1,3,L
+5,4,1,4,L
+5,4,1,5,L
+5,4,2,1,L
+5,4,2,2,L
+5,4,2,3,L
+5,4,2,4,L
+5,4,2,5,L
+5,4,3,1,L
+5,4,3,2,L
+5,4,3,3,L
+5,4,3,4,L
+5,4,3,5,L
+5,4,4,1,L
+5,4,4,2,L
+5,4,4,3,L
+5,4,4,4,L
+5,4,4,5,B
+5,4,5,1,L
+5,4,5,2,L
+5,4,5,3,L
+5,4,5,4,B
+5,4,5,5,R
+5,5,1,1,L
+5,5,1,2,L
+5,5,1,3,L
+5,5,1,4,L
+5,5,1,5,L
+5,5,2,1,L
+5,5,2,2,L
+5,5,2,3,L
+5,5,2,4,L
+5,5,2,5,L
+5,5,3,1,L
+5,5,3,2,L
+5,5,3,3,L
+5,5,3,4,L
+5,5,3,5,L
+5,5,4,1,L
+5,5,4,2,L
+5,5,4,3,L
+5,5,4,4,L
+5,5,4,5,L
+5,5,5,1,L
+5,5,5,2,L
+5,5,5,3,L
+5,5,5,4,L
+5,5,5,5,B
+%
+%
+%

iris.csv

@@ -0,0 +1,151 @@
+5.1,3.5,1.4,0.2,Iris-setosa
+4.9,3.0,1.4,0.2,Iris-setosa
+4.7,3.2,1.3,0.2,Iris-setosa
+4.6,3.1,1.5,0.2,Iris-setosa
+5.0,3.6,1.4,0.2,Iris-setosa
+5.4,3.9,1.7,0.4,Iris-setosa
+4.6,3.4,1.4,0.3,Iris-setosa
+5.0,3.4,1.5,0.2,Iris-setosa
+4.4,2.9,1.4,0.2,Iris-setosa
+4.9,3.1,1.5,0.1,Iris-setosa
+5.4,3.7,1.5,0.2,Iris-setosa
+4.8,3.4,1.6,0.2,Iris-setosa
+4.8,3.0,1.4,0.1,Iris-setosa
+4.3,3.0,1.1,0.1,Iris-setosa
+5.8,4.0,1.2,0.2,Iris-setosa
+5.7,4.4,1.5,0.4,Iris-setosa
+5.4,3.9,1.3,0.4,Iris-setosa
+5.1,3.5,1.4,0.3,Iris-setosa
+5.7,3.8,1.7,0.3,Iris-setosa
+5.1,3.8,1.5,0.3,Iris-setosa
+5.4,3.4,1.7,0.2,Iris-setosa
+5.1,3.7,1.5,0.4,Iris-setosa
+4.6,3.6,1.0,0.2,Iris-setosa
+5.1,3.3,1.7,0.5,Iris-setosa
+4.8,3.4,1.9,0.2,Iris-setosa
+5.0,3.0,1.6,0.2,Iris-setosa
+5.0,3.4,1.6,0.4,Iris-setosa
+5.2,3.5,1.5,0.2,Iris-setosa
+5.2,3.4,1.4,0.2,Iris-setosa
+4.7,3.2,1.6,0.2,Iris-setosa
+4.8,3.1,1.6,0.2,Iris-setosa
+5.4,3.4,1.5,0.4,Iris-setosa
+5.2,4.1,1.5,0.1,Iris-setosa
+5.5,4.2,1.4,0.2,Iris-setosa
+4.9,3.1,1.5,0.1,Iris-setosa
+5.0,3.2,1.2,0.2,Iris-setosa
+5.5,3.5,1.3,0.2,Iris-setosa
+4.9,3.1,1.5,0.1,Iris-setosa
+4.4,3.0,1.3,0.2,Iris-setosa
+5.1,3.4,1.5,0.2,Iris-setosa
+5.0,3.5,1.3,0.3,Iris-setosa
+4.5,2.3,1.3,0.3,Iris-setosa
+4.4,3.2,1.3,0.2,Iris-setosa
+5.0,3.5,1.6,0.6,Iris-setosa
+5.1,3.8,1.9,0.4,Iris-setosa
+4.8,3.0,1.4,0.3,Iris-setosa
+5.1,3.8,1.6,0.2,Iris-setosa
+4.6,3.2,1.4,0.2,Iris-setosa
+5.3,3.7,1.5,0.2,Iris-setosa
+5.0,3.3,1.4,0.2,Iris-setosa
+7.0,3.2,4.7,1.4,Iris-versicolor
+6.4,3.2,4.5,1.5,Iris-versicolor
+6.9,3.1,4.9,1.5,Iris-versicolor
+5.5,2.3,4.0,1.3,Iris-versicolor
+6.5,2.8,4.6,1.5,Iris-versicolor
+5.7,2.8,4.5,1.3,Iris-versicolor
+6.3,3.3,4.7,1.6,Iris-versicolor
+4.9,2.4,3.3,1.0,Iris-versicolor
+6.6,2.9,4.6,1.3,Iris-versicolor
+5.2,2.7,3.9,1.4,Iris-versicolor
+5.0,2.0,3.5,1.0,Iris-versicolor
+5.9,3.0,4.2,1.5,Iris-versicolor
+6.0,2.2,4.0,1.0,Iris-versicolor
+6.1,2.9,4.7,1.4,Iris-versicolor
+5.6,2.9,3.6,1.3,Iris-versicolor
+6.7,3.1,4.4,1.4,Iris-versicolor
+5.6,3.0,4.5,1.5,Iris-versicolor
+5.8,2.7,4.1,1.0,Iris-versicolor
+6.2,2.2,4.5,1.5,Iris-versicolor
+5.6,2.5,3.9,1.1,Iris-versicolor
+5.9,3.2,4.8,1.8,Iris-versicolor
+6.1,2.8,4.0,1.3,Iris-versicolor
+6.3,2.5,4.9,1.5,Iris-versicolor
+6.1,2.8,4.7,1.2,Iris-versicolor
+6.4,2.9,4.3,1.3,Iris-versicolor
+6.6,3.0,4.4,1.4,Iris-versicolor
+6.8,2.8,4.8,1.4,Iris-versicolor
+6.7,3.0,5.0,1.7,Iris-versicolor
+6.0,2.9,4.5,1.5,Iris-versicolor
+5.7,2.6,3.5,1.0,Iris-versicolor
+5.5,2.4,3.8,1.1,Iris-versicolor
+5.5,2.4,3.7,1.0,Iris-versicolor
+5.8,2.7,3.9,1.2,Iris-versicolor
+6.0,2.7,5.1,1.6,Iris-versicolor
+5.4,3.0,4.5,1.5,Iris-versicolor
+6.0,3.4,4.5,1.6,Iris-versicolor
+6.7,3.1,4.7,1.5,Iris-versicolor
+6.3,2.3,4.4,1.3,Iris-versicolor
+5.6,3.0,4.1,1.3,Iris-versicolor
+5.5,2.5,4.0,1.3,Iris-versicolor
+5.5,2.6,4.4,1.2,Iris-versicolor
+6.1,3.0,4.6,1.4,Iris-versicolor
+5.8,2.6,4.0,1.2,Iris-versicolor
+5.0,2.3,3.3,1.0,Iris-versicolor
+5.6,2.7,4.2,1.3,Iris-versicolor
+5.7,3.0,4.2,1.2,Iris-versicolor
+5.7,2.9,4.2,1.3,Iris-versicolor
+6.2,2.9,4.3,1.3,Iris-versicolor
+5.1,2.5,3.0,1.1,Iris-versicolor
+5.7,2.8,4.1,1.3,Iris-versicolor
+6.3,3.3,6.0,2.5,Iris-virginica
+5.8,2.7,5.1,1.9,Iris-virginica
+7.1,3.0,5.9,2.1,Iris-virginica
+6.3,2.9,5.6,1.8,Iris-virginica
+6.5,3.0,5.8,2.2,Iris-virginica
+7.6,3.0,6.6,2.1,Iris-virginica
+4.9,2.5,4.5,1.7,Iris-virginica
+7.3,2.9,6.3,1.8,Iris-virginica
+6.7,2.5,5.8,1.8,Iris-virginica
+7.2,3.6,6.1,2.5,Iris-virginica
+6.5,3.2,5.1,2.0,Iris-virginica
+6.4,2.7,5.3,1.9,Iris-virginica
+6.8,3.0,5.5,2.1,Iris-virginica
+5.7,2.5,5.0,2.0,Iris-virginica
+5.8,2.8,5.1,2.4,Iris-virginica
+6.4,3.2,5.3,2.3,Iris-virginica
+6.5,3.0,5.5,1.8,Iris-virginica
+7.7,3.8,6.7,2.2,Iris-virginica
+7.7,2.6,6.9,2.3,Iris-virginica
+6.0,2.2,5.0,1.5,Iris-virginica
+6.9,3.2,5.7,2.3,Iris-virginica
+5.6,2.8,4.9,2.0,Iris-virginica
+7.7,2.8,6.7,2.0,Iris-virginica
+6.3,2.7,4.9,1.8,Iris-virginica
+6.7,3.3,5.7,2.1,Iris-virginica
+7.2,3.2,6.0,1.8,Iris-virginica
+6.2,2.8,4.8,1.8,Iris-virginica
+6.1,3.0,4.9,1.8,Iris-virginica
+6.4,2.8,5.6,2.1,Iris-virginica
+7.2,3.0,5.8,1.6,Iris-virginica
+7.4,2.8,6.1,1.9,Iris-virginica
+7.9,3.8,6.4,2.0,Iris-virginica
+6.4,2.8,5.6,2.2,Iris-virginica
+6.3,2.8,5.1,1.5,Iris-virginica
+6.1,2.6,5.6,1.4,Iris-virginica
+7.7,3.0,6.1,2.3,Iris-virginica
+6.3,3.4,5.6,2.4,Iris-virginica
+6.4,3.1,5.5,1.8,Iris-virginica
+6.0,3.0,4.8,1.8,Iris-virginica
+6.9,3.1,5.4,2.1,Iris-virginica
+6.7,3.1,5.6,2.4,Iris-virginica
+6.9,3.1,5.1,2.3,Iris-virginica
+5.8,2.7,5.1,1.9,Iris-virginica
+6.8,3.2,5.9,2.3,Iris-virginica
+6.7,3.3,5.7,2.5,Iris-virginica
+6.7,3.0,5.2,2.3,Iris-virginica
+6.3,2.5,5.0,1.9,Iris-virginica
+6.5,3.0,5.2,2.0,Iris-virginica
+6.2,3.4,5.4,2.3,Iris-virginica
+5.9,3.0,5.1,1.8,Iris-virginica
+

requirements.txt

@@ -0,0 +1,16 @@
+# Please use python V 3.7 to be compatible with all packages
+gpytorch==1.5.0
+torch==1.9.0
+scikit-learn==0.24.2
+pyyaml==5.4.1
+seaborn==0.11.2
+xgboost==1.4.0
+tqdm==4.62.1
+numpy==1.21.2
+openml==0.12.2
+catboost==0.26.1
+auto-sklearn==0.14.5
+hyperopt==0.2.5
+configspace==0.4.21
+# autogluon==0.4.0
+gradio==3.1.1