Fix torch segfault in colab example

examples/pysr_demo.ipynb

@@ -152,11 +152,6 @@
         "import numpy as np\n",
         "from matplotlib import pyplot as plt\n",
         "from pysr import PySRRegressor\n",
-        "import torch\n",
-        "from torch import nn, optim\n",
-        "from torch.nn import functional as F\n",
-        "from torch.utils.data import DataLoader, TensorDataset\n",
-        "import pytorch_lightning as pl\n",
         "from sklearn.model_selection import train_test_split"
       ]
     },
@@ -232,8 +227,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "p4PSrO-NK1Wa",
-        "scrolled": true
+        "id": "p4PSrO-NK1Wa"
       },
       "outputs": [],
       "source": [
@@ -412,8 +406,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "PoEkpvYuGUdy",
-        "scrolled": true
+        "id": "PoEkpvYuGUdy"
       },
       "outputs": [],
       "source": [
@@ -606,8 +599,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "a07K3KUjOxcp",
-        "scrolled": true
+        "id": "a07K3KUjOxcp"
       },
       "outputs": [],
       "source": [
@@ -947,8 +939,8 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
+      "id": "ee30bd41",
       "metadata": {},
       "source": [
         "We are all set to go! Let's see if we can find the true relation:"
@@ -1019,10 +1011,13 @@
       },
       "outputs": [],
       "source": [
-        "###### np.random.seed(0)\n",
+        "import numpy as np\n",
+        "\n",
+        "rstate = np.random.RandomState(0)\n",
+        "\n",
         "N = 100000\n",
         "Nt = 10\n",
-        "X = 6 * np.random.rand(N, Nt, 5) - 3\n",
+        "X = 6 * rstate.rand(N, Nt, 5) - 3\n",
         "y_i = X[..., 0] ** 2 + 6 * np.cos(2 * X[..., 2])\n",
         "y = np.sum(y_i, axis=1) / y_i.shape[1]\n",
         "z = y**2\n",
@@ -1055,6 +1050,17 @@
         "Then, we will fit `g` and `f` **separately** using symbolic regression."
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aca54ffa"
+      },
+      "source": [
+        "> **Warning**\n",
+        ">\n",
+        "> We import torch *after* already starting PyJulia. This is required due to interference between their C bindings. If you use torch, and then run PyJulia, you will likely hit a segfault. So keep this in mind for mixed deep learning + PyJulia/PySR workflows."
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -1063,9 +1069,14 @@
       },
       "outputs": [],
       "source": [
-        "hidden = 128\n",
-        "total_steps = 10_000\n",
+        "import torch\n",
+        "from torch import nn, optim\n",
+        "from torch.nn import functional as F\n",
+        "from torch.utils.data import DataLoader, TensorDataset\n",
+        "import pytorch_lightning as pl\n",
         "\n",
+        "hidden = 128\n",
+        "total_steps = 30_000\n",
         "\n",
         "def mlp(size_in, size_out, act=nn.ReLU):\n",
         "    return nn.Sequential(\n",
@@ -1148,13 +1159,14 @@
       },
       "outputs": [],
       "source": [
+        "from multiprocessing import cpu_count\n",
         "Xt = torch.tensor(X).float()\n",
         "zt = torch.tensor(z).float()\n",
         "X_train, X_test, z_train, z_test = train_test_split(Xt, zt, random_state=0)\n",
         "train_set = TensorDataset(X_train, z_train)\n",
-        "train = DataLoader(train_set, batch_size=128, num_workers=2)\n",
+        "train = DataLoader(train_set, batch_size=128, num_workers=cpu_count(), shuffle=True, pin_memory=True)\n",
         "test_set = TensorDataset(X_test, z_test)\n",
-        "test = DataLoader(test_set, batch_size=256, num_workers=2)"
+        "test = DataLoader(test_set, batch_size=256, num_workers=cpu_count(), pin_memory=True)"
       ]
     },
     {
@@ -1207,8 +1219,8 @@
       "outputs": [],
       "source": [
         "trainer = pl.Trainer(\n",
-        "    max_steps=total_steps, accelerator=\"gpu\", devices=1, benchmark=True\n",
-        ")\n"
+        "    max_steps=total_steps, accelerator=\"gpu\", devices=1\n",
+        ")"
       ]
     },
     {
@@ -1262,7 +1274,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "nCCIvvAGuyFi"
@@ -1332,8 +1343,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "51QdHVSkbDhc",
-        "scrolled": true
+        "id": "51QdHVSkbDhc"
       },
       "outputs": [],
       "source": [
@@ -1348,6 +1358,15 @@
         "model.fit(g_input[f_sample_idx], g_output[f_sample_idx])"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1a738a33"
+      },
+      "source": [
+        "If this segfaults, restart the notebook, and run the initial imports and PyJulia part, but skip the PyTorch training. This is because PyTorch's C binding tends to interefere with PyJulia. You can then re-run the `pkl.load` cell to import the data."
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -1380,7 +1399,7 @@
       },
       "outputs": [],
       "source": [
-        "model"
+        "model.equations_[[\"complexity\", \"loss\", \"equation\"]]"
       ]
     },
     {
@@ -1389,7 +1408,7 @@
         "id": "mlU1hidZkgCY"
       },
       "source": [
-        "A neural network can easily undo a linear transform, so this is fine: the network for $f$ will learn to undo the linear transform.\n",
+        "A neural network can easily undo a linear transform (which commutes with the summation), so any affine transform in $g$ is to be expected. The network for $f$ has learned to undo the linear transform.\n",
         "\n",
         "This likely won't find the exact result, but it should find something similar. You may wish to try again but with many more `total_steps` for the neural network (10,000 is quite small!).\n",
         "\n",
@@ -1438,21 +1457,9 @@
     },
     "gpuClass": "standard",
     "kernelspec": {
-      "display_name": "Python (main_ipynb)",
+      "display_name": "Python 3",
       "language": "python",
-      "name": "main_ipynb"
-    },
-    "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.9"
+      "name": "python3"
     }
   },
   "nbformat": 4,