solving-tabular-playground-jan-2021-with-flax.ipynb

{"cells":[{"source":"<a href=\"https://www.kaggle.com/code/yno3fm36xqnnc8/solving-tabular-playground-jan-2021-with-flax?scriptVersionId=141492705\" target=\"_blank\"><img align=\"left\" alt=\"Kaggle\" title=\"Open in Kaggle\" src=\"https://kaggle.com/static/images/open-in-kaggle.svg\"></a>","metadata":{},"cell_type":"markdown"},{"cell_type":"code","execution_count":1,"id":"ce3f5004","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.execute_input":"2023-08-30T21:04:10.70608Z","iopub.status.busy":"2023-08-30T21:04:10.705705Z","iopub.status.idle":"2023-08-30T21:04:13.83179Z","shell.execute_reply":"2023-08-30T21:04:13.83048Z"},"papermill":{"duration":3.140273,"end_time":"2023-08-30T21:04:13.834966","exception":false,"start_time":"2023-08-30T21:04:10.694693","status":"completed"},"tags":[]},"outputs":[{"name":"stderr","output_type":"stream","text":["/opt/conda/lib/python3.10/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.23.5\n","  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"]}],"source":["import pandas as pd\n","import numpy as np\n","\n","import jax\n","import jax.numpy as jnp\n","import flax.linen as nn\n","import flax\n","\n","import matplotlib.pyplot as plt\n","import seaborn as sns\n","\n","from sklearn.model_selection import train_test_split\n","from typing import Sequence, List, Any, Optional"]},{"cell_type":"markdown","id":"489f8441","metadata":{"papermill":{"duration":0.006612,"end_time":"2023-08-30T21:04:13.848764","exception":false,"start_time":"2023-08-30T21:04:13.842152","status":"completed"},"tags":[]},"source":["The goal of this challenge is simply to predict the value of the `target` variable given 14 regressors. This challenge is a regression problem which we'll solve by using a simple MLP, with the `jax` and `flax` libraries helping us out. Growing out of work by the Google Brain team, `flax` is a higher-level library that utilizes `jax` to implement common neural network algorithms."]},{"cell_type":"markdown","id":"a61489ec","metadata":{"papermill":{"duration":0.006518,"end_time":"2023-08-30T21:04:13.862263","exception":false,"start_time":"2023-08-30T21:04:13.855745","status":"completed"},"tags":[]},"source":["First off, let's get the data loaded and check some common summary statistics for the features."]},{"cell_type":"code","execution_count":2,"id":"a6300fe4","metadata":{"execution":{"iopub.execute_input":"2023-08-30T21:04:13.87807Z","iopub.status.busy":"2023-08-30T21:04:13.877365Z","iopub.status.idle":"2023-08-30T21:04:16.351835Z","shell.execute_reply":"2023-08-30T21:04:16.350706Z"},"papermill":{"duration":2.485322,"end_time":"2023-08-30T21:04:16.354335","exception":false,"start_time":"2023-08-30T21:04:13.869013","status":"completed"},"tags":[]},"outputs":[{"data":{"text/html":["<div>\n","<style scoped>\n","    .dataframe tbody tr th:only-of-type {\n","        vertical-align: middle;\n","    }\n","\n","    .dataframe tbody tr th {\n","        vertical-align: top;\n","    }\n","\n","    .dataframe thead th {\n","        text-align: right;\n","    }\n","</style>\n","<table border=\"1\" class=\"dataframe\">\n","  <thead>\n","    <tr style=\"text-align: right;\">\n","      <th></th>\n","      <th>cont1</th>\n","      <th>cont2</th>\n","      <th>cont3</th>\n","      <th>cont4</th>\n","      <th>cont5</th>\n","      <th>cont6</th>\n","      <th>cont7</th>\n","      <th>cont8</th>\n","      <th>cont9</th>\n","      <th>cont10</th>\n","      <th>cont11</th>\n","      <th>cont12</th>\n","      <th>cont13</th>\n","      <th>cont14</th>\n","      <th>target</th>\n","    </tr>\n","  </thead>\n","  <tbody>\n","    <tr>\n","      <th>count</th>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","      <td>300000.000000</td>\n","    </tr>\n","    <tr>\n","      <th>mean</th>\n","      <td>0.506873</td>\n","      <td>0.497898</td>\n","      <td>0.521557</td>\n","      <td>0.515683</td>\n","      <td>0.502022</td>\n","      <td>0.526515</td>\n","      <td>0.487890</td>\n","      <td>0.525163</td>\n","      <td>0.459857</td>\n","      <td>0.520532</td>\n","      <td>0.483926</td>\n","      <td>0.506877</td>\n","      <td>0.553442</td>\n","      <td>0.503713</td>\n","      <td>7.905661</td>\n","    </tr>\n","    <tr>\n","      <th>std</th>\n","      <td>0.203976</td>\n","      <td>0.228159</td>\n","      <td>0.200770</td>\n","      <td>0.233035</td>\n","      <td>0.220701</td>\n","      <td>0.217909</td>\n","      <td>0.181096</td>\n","      <td>0.216221</td>\n","      <td>0.196685</td>\n","      <td>0.201854</td>\n","      <td>0.220082</td>\n","      <td>0.218947</td>\n","      <td>0.229730</td>\n","      <td>0.208238</td>\n","      <td>0.733071</td>\n","    </tr>\n","    <tr>\n","      <th>min</th>\n","      <td>-0.082263</td>\n","      <td>-0.031397</td>\n","      <td>0.020967</td>\n","      <td>0.152761</td>\n","      <td>0.276377</td>\n","      <td>0.066166</td>\n","      <td>-0.097666</td>\n","      <td>0.217260</td>\n","      <td>-0.240604</td>\n","      <td>-0.085046</td>\n","      <td>0.083277</td>\n","      <td>0.088635</td>\n","      <td>0.029950</td>\n","      <td>0.166367</td>\n","      <td>0.000000</td>\n","    </tr>\n","    <tr>\n","      <th>25%</th>\n","      <td>0.343078</td>\n","      <td>0.319170</td>\n","      <td>0.344096</td>\n","      <td>0.294935</td>\n","      <td>0.284108</td>\n","      <td>0.356163</td>\n","      <td>0.346600</td>\n","      <td>0.341486</td>\n","      <td>0.330832</td>\n","      <td>0.375465</td>\n","      <td>0.300474</td>\n","      <td>0.310166</td>\n","      <td>0.350472</td>\n","      <td>0.308673</td>\n","      <td>7.329367</td>\n","    </tr>\n","    <tr>\n","      <th>50%</th>\n","      <td>0.484005</td>\n","      <td>0.553209</td>\n","      <td>0.551471</td>\n","      <td>0.482880</td>\n","      <td>0.451733</td>\n","      <td>0.470988</td>\n","      <td>0.466825</td>\n","      <td>0.483460</td>\n","      <td>0.416843</td>\n","      <td>0.458877</td>\n","      <td>0.441916</td>\n","      <td>0.486599</td>\n","      <td>0.487707</td>\n","      <td>0.431845</td>\n","      <td>7.940571</td>\n","    </tr>\n","    <tr>\n","      <th>75%</th>\n","      <td>0.643789</td>\n","      <td>0.731263</td>\n","      <td>0.648315</td>\n","      <td>0.748705</td>\n","      <td>0.670660</td>\n","      <td>0.694043</td>\n","      <td>0.581292</td>\n","      <td>0.685250</td>\n","      <td>0.575041</td>\n","      <td>0.700292</td>\n","      <td>0.679128</td>\n","      <td>0.694453</td>\n","      <td>0.768479</td>\n","      <td>0.712653</td>\n","      <td>8.470084</td>\n","    </tr>\n","    <tr>\n","      <th>max</th>\n","      <td>1.016227</td>\n","      <td>0.859697</td>\n","      <td>1.006955</td>\n","      <td>1.010402</td>\n","      <td>1.034261</td>\n","      <td>1.043858</td>\n","      <td>1.066167</td>\n","      <td>1.024427</td>\n","      <td>1.004114</td>\n","      <td>1.199951</td>\n","      <td>1.022620</td>\n","      <td>1.049025</td>\n","      <td>0.977845</td>\n","      <td>0.868506</td>\n","      <td>10.267569</td>\n","    </tr>\n","  </tbody>\n","</table>\n","</div>"],"text/plain":["               cont1          cont2          cont3          cont4  \\\n","count  300000.000000  300000.000000  300000.000000  300000.000000   \n","mean        0.506873       0.497898       0.521557       0.515683   \n","std         0.203976       0.228159       0.200770       0.233035   \n","min        -0.082263      -0.031397       0.020967       0.152761   \n","25%         0.343078       0.319170       0.344096       0.294935   \n","50%         0.484005       0.553209       0.551471       0.482880   \n","75%         0.643789       0.731263       0.648315       0.748705   \n","max         1.016227       0.859697       1.006955       1.010402   \n","\n","               cont5          cont6          cont7          cont8  \\\n","count  300000.000000  300000.000000  300000.000000  300000.000000   \n","mean        0.502022       0.526515       0.487890       0.525163   \n","std         0.220701       0.217909       0.181096       0.216221   \n","min         0.276377       0.066166      -0.097666       0.217260   \n","25%         0.284108       0.356163       0.346600       0.341486   \n","50%         0.451733       0.470988       0.466825       0.483460   \n","75%         0.670660       0.694043       0.581292       0.685250   \n","max         1.034261       1.043858       1.066167       1.024427   \n","\n","               cont9         cont10         cont11         cont12  \\\n","count  300000.000000  300000.000000  300000.000000  300000.000000   \n","mean        0.459857       0.520532       0.483926       0.506877   \n","std         0.196685       0.201854       0.220082       0.218947   \n","min        -0.240604      -0.085046       0.083277       0.088635   \n","25%         0.330832       0.375465       0.300474       0.310166   \n","50%         0.416843       0.458877       0.441916       0.486599   \n","75%         0.575041       0.700292       0.679128       0.694453   \n","max         1.004114       1.199951       1.022620       1.049025   \n","\n","              cont13         cont14         target  \n","count  300000.000000  300000.000000  300000.000000  \n","mean        0.553442       0.503713       7.905661  \n","std         0.229730       0.208238       0.733071  \n","min         0.029950       0.166367       0.000000  \n","25%         0.350472       0.308673       7.329367  \n","50%         0.487707       0.431845       7.940571  \n","75%         0.768479       0.712653       8.470084  \n","max         0.977845       0.868506      10.267569  "]},"execution_count":2,"metadata":{},"output_type":"execute_result"}],"source":["train_data = pd.read_csv(\"/kaggle/input/tabular-playground-series-jan-2021/train.csv\").set_index('id')\n","train_data.describe()"]},{"cell_type":"markdown","id":"b8cbbba8","metadata":{"papermill":{"duration":0.007418,"end_time":"2023-08-30T21:04:16.369291","exception":false,"start_time":"2023-08-30T21:04:16.361873","status":"completed"},"tags":[]},"source":["Now we see our first example of a `flax` model, in this case a simple MLP. Its sole argument is `features`, a sequence of integers specifying the size of the dense linear layers in the network."]},{"cell_type":"code","execution_count":3,"id":"8083dda6","metadata":{"execution":{"iopub.execute_input":"2023-08-30T21:04:16.385725Z","iopub.status.busy":"2023-08-30T21:04:16.384743Z","iopub.status.idle":"2023-08-30T21:04:16.391853Z","shell.execute_reply":"2023-08-30T21:04:16.391029Z"},"papermill":{"duration":0.017509,"end_time":"2023-08-30T21:04:16.393907","exception":false,"start_time":"2023-08-30T21:04:16.376398","status":"completed"},"tags":[]},"outputs":[],"source":["class MLP(nn.Module):\n","    features: Sequence[int]\n","            \n","    @nn.compact\n","    def __call__(self, x):\n","        for feat in self.features[:-1]:\n","            x = nn.relu(nn.Dense(feat)(x))\n","        x = nn.Dense(self.features[-1])(x)\n","        return x"]},{"cell_type":"markdown","id":"f9b8d3fb","metadata":{"papermill":{"duration":0.006911,"end_time":"2023-08-30T21:04:16.408169","exception":false,"start_time":"2023-08-30T21:04:16.401258","status":"completed"},"tags":[]},"source":["Now let's actually create a model instance."]},{"cell_type":"code","execution_count":4,"id":"690ec25c","metadata":{"execution":{"iopub.execute_input":"2023-08-30T21:04:16.424483Z","iopub.status.busy":"2023-08-30T21:04:16.42363Z","iopub.status.idle":"2023-08-30T21:04:18.755682Z","shell.execute_reply":"2023-08-30T21:04:18.754396Z"},"papermill":{"duration":2.343301,"end_time":"2023-08-30T21:04:18.758592","exception":false,"start_time":"2023-08-30T21:04:16.415291","status":"completed"},"tags":[]},"outputs":[],"source":["model = MLP(features=[12, 8, 8, 8, 4, 4, 1])\n","\n","variables = model.init(jax.random.PRNGKey(0), jnp.ones((1,14)))\n","vpredict = jax.vmap(model.apply, (None, 0))"]},{"cell_type":"markdown","id":"5840f243","metadata":{"papermill":{"duration":0.007111,"end_time":"2023-08-30T21:04:18.773849","exception":false,"start_time":"2023-08-30T21:04:18.766738","status":"completed"},"tags":[]},"source":["We define the loss function, MSE in this case, and then define its gradient."]},{"cell_type":"code","execution_count":5,"id":"f99d78aa","metadata":{"execution":{"iopub.execute_input":"2023-08-30T21:04:18.790299Z","iopub.status.busy":"2023-08-30T21:04:18.789894Z","iopub.status.idle":"2023-08-30T21:04:18.796223Z","shell.execute_reply":"2023-08-30T21:04:18.795061Z"},"papermill":{"duration":0.017859,"end_time":"2023-08-30T21:04:18.799084","exception":false,"start_time":"2023-08-30T21:04:18.781225","status":"completed"},"tags":[]},"outputs":[],"source":["def loss(variables: flax.core.frozen_dict.FrozenDict, X: jnp.array, y: jnp.array):\n","    return jnp.mean(jnp.square(y - vpredict(variables, X)))\n","loss_grad_fn = jax.value_and_grad(loss)"]},{"cell_type":"markdown","id":"228b7452","metadata":{"papermill":{"duration":0.00733,"end_time":"2023-08-30T21:04:18.813879","exception":false,"start_time":"2023-08-30T21:04:18.806549","status":"completed"},"tags":[]},"source":["Below is the training loop. In this notebook I'm doing this by hand, but in real life it may be better/more convenient to use a library like `optax`."]},{"cell_type":"code","execution_count":6,"id":"ab756a14","metadata":{"execution":{"iopub.execute_input":"2023-08-30T21:04:18.83011Z","iopub.status.busy":"2023-08-30T21:04:18.829717Z","iopub.status.idle":"2023-08-30T21:04:23.180332Z","shell.execute_reply":"2023-08-30T21:04:23.179339Z"},"papermill":{"duration":4.362006,"end_time":"2023-08-30T21:04:23.183134","exception":false,"start_time":"2023-08-30T21:04:18.821128","status":"completed"},"tags":[]},"outputs":[{"name":"stdout","output_type":"stream","text":["62.60427\n","23.943077\n"]}],"source":["learning_rate = 1e-2\n","last_value = None\n","variables = flax.core.frozen_dict.unfreeze(variables)\n","data = train_data.copy().to_numpy()\n","rng = np.random.default_rng()\n","for i in range(1000):\n","    X = rng.choice(data, 1000, axis=0, replace=False)\n","    X_train, y_train = X[:,:-1], X[:,-1]\n","    value, grads = loss_grad_fn(variables, X_train, y_train)\n","    \n","    if last_value is not None and abs(last_value - value) < 0.001 and value < 1.0:\n","        break\n","    \n","    last_value = value\n","    if i % 10 == 0:\n","        print(value)\n","    for key in grads['params']:\n","        for k in grads['params'][key]:\n","            variables['params'][key][k] -= learning_rate * grads['params'][key][k]\n","\n","variables = flax.core.frozen_dict.freeze(variables)"]},{"cell_type":"markdown","id":"b051bcf1","metadata":{"papermill":{"duration":0.007214,"end_time":"2023-08-30T21:04:23.198212","exception":false,"start_time":"2023-08-30T21:04:23.190998","status":"completed"},"tags":[]},"source":["Finally, let's predict the target value for the test dataset and prepare the submission."]},{"cell_type":"code","execution_count":7,"id":"13529f9a","metadata":{"execution":{"iopub.execute_input":"2023-08-30T21:04:23.215747Z","iopub.status.busy":"2023-08-30T21:04:23.214512Z","iopub.status.idle":"2023-08-30T21:04:25.07486Z","shell.execute_reply":"2023-08-30T21:04:25.073637Z"},"papermill":{"duration":1.872029,"end_time":"2023-08-30T21:04:25.077733","exception":false,"start_time":"2023-08-30T21:04:23.205704","status":"completed"},"tags":[]},"outputs":[{"data":{"text/html":["<div>\n","<style scoped>\n","    .dataframe tbody tr th:only-of-type {\n","        vertical-align: middle;\n","    }\n","\n","    .dataframe tbody tr th {\n","        vertical-align: top;\n","    }\n","\n","    .dataframe thead th {\n","        text-align: right;\n","    }\n","</style>\n","<table border=\"1\" class=\"dataframe\">\n","  <thead>\n","    <tr style=\"text-align: right;\">\n","      <th></th>\n","      <th>cont1</th>\n","      <th>cont2</th>\n","      <th>cont3</th>\n","      <th>cont4</th>\n","      <th>cont5</th>\n","      <th>cont6</th>\n","      <th>cont7</th>\n","      <th>cont8</th>\n","      <th>cont9</th>\n","      <th>cont10</th>\n","      <th>cont11</th>\n","      <th>cont12</th>\n","      <th>cont13</th>\n","      <th>cont14</th>\n","      <th>target</th>\n","    </tr>\n","    <tr>\n","      <th>id</th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","      <th></th>\n","    </tr>\n","  </thead>\n","  <tbody>\n","    <tr>\n","      <th>0</th>\n","      <td>0.353600</td>\n","      <td>0.738780</td>\n","      <td>0.600939</td>\n","      <td>0.293377</td>\n","      <td>0.285691</td>\n","      <td>0.458006</td>\n","      <td>0.620704</td>\n","      <td>0.422249</td>\n","      <td>0.369203</td>\n","      <td>0.435727</td>\n","      <td>0.550540</td>\n","      <td>0.699134</td>\n","      <td>0.286864</td>\n","      <td>0.364515</td>\n","      <td>7.427158</td>\n","    </tr>\n","    <tr>\n","      <th>2</th>\n","      <td>0.907222</td>\n","      <td>0.189756</td>\n","      <td>0.215531</td>\n","      <td>0.869915</td>\n","      <td>0.301333</td>\n","      <td>0.528958</td>\n","      <td>0.390351</td>\n","      <td>0.521112</td>\n","      <td>0.794779</td>\n","      <td>0.798580</td>\n","      <td>0.446475</td>\n","      <td>0.449037</td>\n","      <td>0.916964</td>\n","      <td>0.513002</td>\n","      <td>7.240384</td>\n","    </tr>\n","    <tr>\n","      <th>6</th>\n","      <td>0.179287</td>\n","      <td>0.355353</td>\n","      <td>0.623972</td>\n","      <td>0.437812</td>\n","      <td>0.282476</td>\n","      <td>0.320826</td>\n","      <td>0.386789</td>\n","      <td>0.776422</td>\n","      <td>0.222268</td>\n","      <td>0.229102</td>\n","      <td>0.211913</td>\n","      <td>0.222651</td>\n","      <td>0.327164</td>\n","      <td>0.827941</td>\n","      <td>6.783124</td>\n","    </tr>\n","    <tr>\n","      <th>7</th>\n","      <td>0.359385</td>\n","      <td>0.181049</td>\n","      <td>0.551368</td>\n","      <td>0.206386</td>\n","      <td>0.280763</td>\n","      <td>0.482076</td>\n","      <td>0.506677</td>\n","      <td>0.362793</td>\n","      <td>0.379737</td>\n","      <td>0.345686</td>\n","      <td>0.445276</td>\n","      <td>0.518485</td>\n","      <td>0.299028</td>\n","      <td>0.598166</td>\n","      <td>6.783523</td>\n","    </tr>\n","    <tr>\n","      <th>10</th>\n","      <td>0.335791</td>\n","      <td>0.682607</td>\n","      <td>0.676481</td>\n","      <td>0.219465</td>\n","      <td>0.282861</td>\n","      <td>0.581721</td>\n","      <td>0.748639</td>\n","      <td>0.350158</td>\n","      <td>0.448915</td>\n","      <td>0.506878</td>\n","      <td>0.817721</td>\n","      <td>0.805895</td>\n","      <td>0.790591</td>\n","      <td>0.249275</td>\n","      <td>7.811783</td>\n","    </tr>\n","    <tr>\n","      <th>...</th>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","      <td>...</td>\n","    </tr>\n","    <tr>\n","      <th>499984</th>\n","      <td>0.353856</td>\n","      <td>0.677578</td>\n","      <td>0.550852</td>\n","      <td>0.869612</td>\n","      <td>0.957635</td>\n","      <td>0.255054</td>\n","      <td>0.289138</td>\n","      <td>0.635979</td>\n","      <td>0.271399</td>\n","      <td>0.282455</td>\n","      <td>0.217169</td>\n","      <td>0.219088</td>\n","      <td>0.373261</td>\n","      <td>0.272479</td>\n","      <td>7.257768</td>\n","    </tr>\n","    <tr>\n","      <th>499985</th>\n","      <td>0.243209</td>\n","      <td>0.135627</td>\n","      <td>0.218393</td>\n","      <td>0.792798</td>\n","      <td>0.547639</td>\n","      <td>0.433520</td>\n","      <td>0.549540</td>\n","      <td>0.650107</td>\n","      <td>0.453787</td>\n","      <td>0.459689</td>\n","      <td>0.450424</td>\n","      <td>0.511176</td>\n","      <td>0.318334</td>\n","      <td>0.395747</td>\n","      <td>7.143597</td>\n","    </tr>\n","    <tr>\n","      <th>499987</th>\n","      <td>0.506973</td>\n","      <td>0.683893</td>\n","      <td>0.533434</td>\n","      <td>0.192957</td>\n","      <td>0.314381</td>\n","      <td>0.358604</td>\n","      <td>0.554455</td>\n","      <td>0.267105</td>\n","      <td>0.396101</td>\n","      <td>0.445390</td>\n","      <td>0.382656</td>\n","      <td>0.397978</td>\n","      <td>0.381235</td>\n","      <td>0.369464</td>\n","      <td>7.238574</td>\n","    </tr>\n","    <tr>\n","      <th>499988</th>\n","      <td>0.347870</td>\n","      <td>0.553112</td>\n","      <td>0.495284</td>\n","      <td>0.861500</td>\n","      <td>0.816914</td>\n","      <td>0.298478</td>\n","      <td>0.275964</td>\n","      <td>0.265841</td>\n","      <td>0.334250</td>\n","      <td>0.252635</td>\n","      <td>0.213589</td>\n","      <td>0.285223</td>\n","      <td>0.336772</td>\n","      <td>0.388505</td>\n","      <td>7.049764</td>\n","    </tr>\n","    <tr>\n","      <th>499990</th>\n","      <td>0.688188</td>\n","      <td>0.358328</td>\n","      <td>0.644511</td>\n","      <td>0.766354</td>\n","      <td>0.747591</td>\n","      <td>0.727329</td>\n","      <td>0.468329</td>\n","      <td>0.576711</td>\n","      <td>0.453124</td>\n","      <td>0.381600</td>\n","      <td>0.299126</td>\n","      <td>0.499627</td>\n","      <td>0.691078</td>\n","      <td>0.668109</td>\n","      <td>7.346300</td>\n","    </tr>\n","  </tbody>\n","</table>\n","<p>200000 rows × 15 columns</p>\n","</div>"],"text/plain":["           cont1     cont2     cont3     cont4     cont5     cont6     cont7  \\\n","id                                                                             \n","0       0.353600  0.738780  0.600939  0.293377  0.285691  0.458006  0.620704   \n","2       0.907222  0.189756  0.215531  0.869915  0.301333  0.528958  0.390351   \n","6       0.179287  0.355353  0.623972  0.437812  0.282476  0.320826  0.386789   \n","7       0.359385  0.181049  0.551368  0.206386  0.280763  0.482076  0.506677   \n","10      0.335791  0.682607  0.676481  0.219465  0.282861  0.581721  0.748639   \n","...          ...       ...       ...       ...       ...       ...       ...   \n","499984  0.353856  0.677578  0.550852  0.869612  0.957635  0.255054  0.289138   \n","499985  0.243209  0.135627  0.218393  0.792798  0.547639  0.433520  0.549540   \n","499987  0.506973  0.683893  0.533434  0.192957  0.314381  0.358604  0.554455   \n","499988  0.347870  0.553112  0.495284  0.861500  0.816914  0.298478  0.275964   \n","499990  0.688188  0.358328  0.644511  0.766354  0.747591  0.727329  0.468329   \n","\n","           cont8     cont9    cont10    cont11    cont12    cont13    cont14  \\\n","id                                                                             \n","0       0.422249  0.369203  0.435727  0.550540  0.699134  0.286864  0.364515   \n","2       0.521112  0.794779  0.798580  0.446475  0.449037  0.916964  0.513002   \n","6       0.776422  0.222268  0.229102  0.211913  0.222651  0.327164  0.827941   \n","7       0.362793  0.379737  0.345686  0.445276  0.518485  0.299028  0.598166   \n","10      0.350158  0.448915  0.506878  0.817721  0.805895  0.790591  0.249275   \n","...          ...       ...       ...       ...       ...       ...       ...   \n","499984  0.635979  0.271399  0.282455  0.217169  0.219088  0.373261  0.272479   \n","499985  0.650107  0.453787  0.459689  0.450424  0.511176  0.318334  0.395747   \n","499987  0.267105  0.396101  0.445390  0.382656  0.397978  0.381235  0.369464   \n","499988  0.265841  0.334250  0.252635  0.213589  0.285223  0.336772  0.388505   \n","499990  0.576711  0.453124  0.381600  0.299126  0.499627  0.691078  0.668109   \n","\n","          target  \n","id                \n","0       7.427158  \n","2       7.240384  \n","6       6.783124  \n","7       6.783523  \n","10      7.811783  \n","...          ...  \n","499984  7.257768  \n","499985  7.143597  \n","499987  7.238574  \n","499988  7.049764  \n","499990  7.346300  \n","\n","[200000 rows x 15 columns]"]},"execution_count":7,"metadata":{},"output_type":"execute_result"}],"source":["test_data = pd.read_csv(\"/kaggle/input/tabular-playground-series-jan-2021/test.csv\").set_index('id')\n","test_data['target'] = vpredict(variables, jnp.array(test_data.to_numpy()))\n","\n","test_data"]},{"cell_type":"code","execution_count":8,"id":"0e8cdf61","metadata":{"execution":{"iopub.execute_input":"2023-08-30T21:04:25.095758Z","iopub.status.busy":"2023-08-30T21:04:25.095356Z","iopub.status.idle":"2023-08-30T21:04:25.82036Z","shell.execute_reply":"2023-08-30T21:04:25.818991Z"},"papermill":{"duration":0.737163,"end_time":"2023-08-30T21:04:25.82323","exception":false,"start_time":"2023-08-30T21:04:25.086067","status":"completed"},"tags":[]},"outputs":[],"source":["test_data[['target']].to_csv(\"submission.csv\")"]}],"kernelspec":{"display_name":"Python 3","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.6.4"},"metadata":{"kernelspec":{"display_name":"Python 3","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.10.12"},"papermill":{"default_parameters":{},"duration":27.263963,"end_time":"2023-08-30T21:04:27.059812","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2023-08-30T21:03:59.795849","version":"2.4.0"}},"nbformat":4,"nbformat_minor":5}