diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index e238de2c8..8be9ae637 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -83,7 +83,6 @@ jobs:
         shell: bash -l {0}
     strategy:
       matrix:
-        julia-version: ['1.7.1']
         python-version: ['3.9']
         os: ['ubuntu-latest']
     
@@ -108,7 +107,7 @@ jobs:
       - name: "Cache Julia"
         uses: julia-actions/cache@v1
         with:
-          cache-name: ${{ matrix.os }}-conda-${{ matrix.julia-version }}-${{ matrix.python-version }}
+          cache-name: ${{ matrix.os }}-conda-${{ matrix.python-version }}
           cache-packages: false
       - name: "Install PySR"
         run: |
diff --git a/.github/workflows/CI_Windows.yml b/.github/workflows/CI_Windows.yml
index 08b3cb5c7..53dff82ec 100644
--- a/.github/workflows/CI_Windows.yml
+++ b/.github/workflows/CI_Windows.yml
@@ -29,7 +29,7 @@ jobs:
         shell: bash
     strategy:
       matrix:
-        julia-version: ['1.6', '1.8.2']
+        julia-version: ['1.8.2']
         python-version: ['3.9']
         os: [windows-latest]
     
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 93a6fde92..42c5975a5 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -27,7 +27,7 @@ jobs:
           python-version: 3.9
           cache: pip
       - name: "Install packages for docs building"
-        run: pip install mkdocs-material mkdocs-autorefs 'mkdocstrings[python]' docstring_parser
+        run: pip install -r docs/requirements.txt
       - name: "Install PySR"
         run: pip install -e .
       - name: "Build API docs"
diff --git a/.github/workflows/pypi_deploy.yml b/.github/workflows/pypi_deploy.yml
index 52b9d92dc..9124e35a3 100644
--- a/.github/workflows/pypi_deploy.yml
+++ b/.github/workflows/pypi_deploy.yml
@@ -3,17 +3,12 @@ on:
   push:
     tags:
       - 'v*.*.*'
+  workflow_dispatch:
 
 jobs:
   pypi:
     runs-on: ubuntu-latest
     steps:
-      - name: Wait for tests to pass
-        uses: lewagon/wait-on-check-action@v1.2.0
-        with:
-          ref: ${{ github.ref }}
-          check-name: 'Linux'
-          repo-token: ${{ secrets.GITHUB_TOKEN }}
       - name: "Checkout"
         uses: actions/checkout@v3
       - name: "Set up Python"
diff --git a/README.md b/README.md
index 6a242337c..e6adb436f 100644
--- a/README.md
+++ b/README.md
@@ -87,32 +87,37 @@ If none of these folders contain your Julia binary, then you need to add Julia's
 
 # Introduction
 
-Let's create a PySR example. First, let's import
-numpy to generate some test data:
+You might wish to try the interactive tutorial [here](https://colab.research.google.com/github/MilesCranmer/PySR/blob/master/examples/pysr_demo.ipynb), which uses the notebook in `examples/pysr_demo.ipynb`.
+
+In practice, I highly recommend using IPython rather than Jupyter, as the printing is much nicer.
+Below is a quick demo here which you can paste into a Python runtime.
+First, let's import numpy to generate some test data:
+
 ```python
 import numpy as np
 
 X = 2 * np.random.randn(100, 5)
 y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5
 ```
+
 We have created a dataset with 100 datapoints, with 5 features each.
 The relation we wish to model is $2.5382 \cos(x_3) + x_0^2 - 0.5$.
 
 Now, let's create a PySR model and train it.
 PySR's main interface is in the style of scikit-learn:
+
 ```python
 from pysr import PySRRegressor
 
 model = PySRRegressor(
-    model_selection="best",  # Result is mix of simplicity+accuracy
-    niterations=40,
+    niterations=40,  # < Increase me for better results
     binary_operators=["+", "*"],
     unary_operators=[
         "cos",
         "exp",
         "sin",
         "inv(x) = 1/x",
-	# ^ Custom operator (julia syntax)
+        # ^ Custom operator (julia syntax)
     ],
     extra_sympy_mappings={"inv": lambda x: 1 / x},
     # ^ Define operator for SymPy as well
@@ -120,25 +125,32 @@ model = PySRRegressor(
     # ^ Custom loss function (julia syntax)
 )
 ```
+
 This will set up the model for 40 iterations of the search code, which contains hundreds of thousands of mutations and equation evaluations.
 
 Let's train this model on our dataset:
+
 ```python
 model.fit(X, y)
 ```
+
 Internally, this launches a Julia process which will do a multithreaded search for equations to fit the dataset.
 
 Equations will be printed during training, and once you are satisfied, you may 
 quit early by hitting 'q' and then \<enter\>.
 
 After the model has been fit, you can run `model.predict(X)`
-to see the predictions on a given dataset.
+to see the predictions on a given dataset using the automatically-selected expression,
+or, for example, `model.predict(X, 3)` to see the predictions of the 3rd equation.
 
 You may run:
+
 ```python
 print(model)
 ```
+
 to print the learned equations:
+
 ```python
 PySRRegressor.equations_ = [
 	   pick     score                                           equation       loss  complexity
@@ -150,6 +162,7 @@ PySRRegressor.equations_ = [
 	5  >>>>       inf  (((cos(x3) + -0.19699033) * 2.5382123) + (x0 *...   0.000000          10
 ]
 ```
+
 This arrow in the `pick` column indicates which equation is currently selected by your
 `model_selection` strategy for prediction.
 (You may change `model_selection` after `.fit(X, y)` as well.)
@@ -165,6 +178,7 @@ This will cause problems if significant changes are made to the search parameter
 You will notice that PySR will save two files: `hall_of_fame...csv` and `hall_of_fame...pkl`.
 The csv file is a list of equations and their losses, and the pkl file is a saved state of the model.
 You may load the model from the `pkl` file with:
+
 ```python
 model = PySRRegressor.from_file("hall_of_fame.2022-08-10_100832.281.pkl")
 ``` 
@@ -254,22 +268,25 @@ model = PySRRegressor(
 )
 ```
 
-
 # Docker
 
 You can also test out PySR in Docker, without
 installing it locally, by running the following command in
 the root directory of this repo:
+
 ```bash
 docker build -t pysr .
 ```
+
 This builds an image called `pysr` for your system's architecture,
 which also contains IPython.
 
 You can then run this with:
+
 ```bash
 docker run -it --rm -v "$PWD:/data" pysr ipython
 ```
+
 which will link the current directory to the container's `/data` directory
 and then launch ipython.
 
diff --git a/docs/README.md b/docs/README.md
new file mode 100644
index 000000000..34f4e7d85
--- /dev/null
+++ b/docs/README.md
@@ -0,0 +1,8 @@
+# PySR Documentation
+
+## Building locally
+
+1. In the base directory, run `pip install -r docs/requirements.txt`.
+2. Install PySR in editable mode: `pip install -e .`.
+3. Build doc source with `cd docs && ./gen_docs.sh && cd ..`.
+4. Create and serve docs with mkdocs: `mkdocs serve -w pysr`.
diff --git a/docs/examples.md b/docs/examples.md
index 2b35b6294..14d3321fb 100644
--- a/docs/examples.md
+++ b/docs/examples.md
@@ -249,9 +249,11 @@ Next, let's use this list of primes to create a dataset of $x, y$ pairs:
 import numpy as np
 
 X = np.random.randint(0, 100, 100)[:, None]
-y = [primes[3*X[i, 0] + 1] - 5 for i in range(100)]
+y = [primes[3*X[i, 0] + 1] - 5 + np.random.randn()*0.001 for i in range(100)]
 ```
 
+Note that we have also added a tiny bit of noise to the dataset.
+
 Finally, let's create a PySR model, and pass the custom operator. We also need to define the sympy equivalent, which we can leave as a placeholder for now:
 
 ```python
@@ -264,7 +266,7 @@ class sympy_p(sympy.Function):
 model = PySRRegressor(
     binary_operators=["+", "-", "*", "/"],
     unary_operators=["p"],
-    niterations=1000,
+    niterations=100,
     extra_sympy_mappings={"p": sympy_p}
 )
 ```
@@ -276,10 +278,10 @@ model.fit(X, y)
 ```
 
 if all works out, you should be able to see the true relation (note that the constant offset might not be exactly 1, since it is allowed to round to the nearest integer).
-You can get the sympy version of the last row with:
+You can get the sympy version of the best equation with:
 
 ```python
-model.sympy(index=-1)
+model.sympy()
 ```
 
 ## 8. Additional features
diff --git a/docs/interactive-docs.md b/docs/interactive-docs.md
new file mode 100644
index 000000000..3c87a5107
--- /dev/null
+++ b/docs/interactive-docs.md
@@ -0,0 +1,11 @@
+# Interactive Reference ⭐
+
+<!-- Display content from `astroautomata.com/pysr_interactive` -->
+
+The following docs are interactive, and, based on your selections,
+will create a snippet of Python code at the bottom which you can execute locally.
+Clicking on each parameter's name will display a description.
+Note that this is an incomplete list of options; for the full list,
+see the [API Reference](api.md).
+
+<embed src="https://astroautomata.com/pysr_interactive" width="100%" height="2000px" />
diff --git a/docs/requirements.txt b/docs/requirements.txt
new file mode 100644
index 000000000..29381320e
--- /dev/null
+++ b/docs/requirements.txt
@@ -0,0 +1,4 @@
+mkdocs-material
+mkdocs-autorefs
+mkdocstrings[python]
+docstring_parser
\ No newline at end of file
diff --git a/examples/pysr_demo.ipynb b/examples/pysr_demo.ipynb
index 337eb6ac0..68271b1aa 100644
--- a/examples/pysr_demo.ipynb
+++ b/examples/pysr_demo.ipynb
@@ -40,11 +40,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "GIeFXS0F0zww",
-        "outputId": "f25272ac-a660-42fd-d739-82778e6d7415"
+        "id": "GIeFXS0F0zww"
       },
       "outputs": [],
       "source": [
@@ -52,9 +48,8 @@
         "set -e\n",
         "\n",
         "#---------------------------------------------------#\n",
-        "JULIA_VERSION=\"1.7.2\"\n",
-        "JULIA_PACKAGES=\"PyCall SymbolicRegression\"\n",
-        "JULIA_NUM_THREADS=4\n",
+        "JULIA_VERSION=\"1.8.5\"\n",
+        "export JULIA_PKG_PRECOMPILE_AUTO=0\n",
         "#---------------------------------------------------#\n",
         "\n",
         "if [ -z `which julia` ]; then\n",
@@ -67,12 +62,10 @@
         "  tar -x -f /tmp/julia.tar.gz -C /usr/local --strip-components 1\n",
         "  rm /tmp/julia.tar.gz\n",
         "\n",
-        "  for PKG in `echo $JULIA_PACKAGES`; do\n",
-        "    echo \"Installing Julia package $PKG...\"\n",
-        "    julia -e 'using Pkg; pkg\"add '$PKG'; precompile;\"'\n",
-        "  done\n",
-        "  \n",
+        "  echo \"Installing PyCall.jl...\"\n",
+        "  julia -e 'using Pkg; Pkg.add(\"PyCall\"); Pkg.build(\"PyCall\")'\n",
         "  julia -e 'println(\"Success\")'\n",
+        "\n",
         "fi"
       ]
     },
@@ -98,7 +91,9 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "etTMEV0wDqld"
+      },
       "source": [
         "The following step is not normally required, but colab's printing is non-standard and we need to manually set it up PyJulia:\n"
       ]
@@ -113,11 +108,11 @@
       "source": [
         "from julia import Julia\n",
         "\n",
-        "julia = Julia(compiled_modules=False)\n",
+        "julia = Julia(compiled_modules=False, threads='auto')\n",
         "from julia import Main\n",
         "from julia.tools import redirect_output_streams\n",
         "\n",
-        "redirect_output_streams()\n"
+        "redirect_output_streams()"
       ]
     },
     {
@@ -126,7 +121,7 @@
         "id": "6u2WhbVhht-G"
       },
       "source": [
-        "Let's install the backend of PySR, and all required libraries. We will also precompile them so they are faster at startup.\n",
+        "Let's install the backend of PySR, and all required libraries.\n",
         "\n",
         "**(This may take some time)**"
       ]
@@ -135,17 +130,14 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "J-0QbxyK1_51",
-        "outputId": "3742d548-3574-4739-80b1-ee9e906e4b57"
+        "id": "J-0QbxyK1_51"
       },
       "outputs": [],
       "source": [
         "import pysr\n",
         "\n",
-        "pysr.install()\n"
+        "# We don't precompile in colab because compiled modules are incompatible static Python libraries:\n",
+        "pysr.install(precompile=False)"
       ]
     },
     {
@@ -160,12 +152,7 @@
         "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\n"
+        "from sklearn.model_selection import train_test_split"
       ]
     },
     {
@@ -199,7 +186,7 @@
         "# Dataset\n",
         "np.random.seed(0)\n",
         "X = 2 * np.random.randn(100, 5)\n",
-        "y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 2\n"
+        "y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 2"
       ]
     },
     {
@@ -221,9 +208,8 @@
       "source": [
         "default_pysr_params = dict(\n",
         "    populations=30,\n",
-        "    procs=4,\n",
         "    model_selection=\"best\",\n",
-        ")\n"
+        ")"
       ]
     },
     {
@@ -241,12 +227,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "p4PSrO-NK1Wa",
-        "outputId": "474b4780-5d94-4795-88b9-225030b17abe",
-        "scrolled": true
+        "id": "p4PSrO-NK1Wa"
       },
       "outputs": [],
       "source": [
@@ -258,7 +239,7 @@
         "    **default_pysr_params\n",
         ")\n",
         "\n",
-        "model.fit(X, y)\n"
+        "model.fit(X, y)"
       ]
     },
     {
@@ -274,15 +255,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "4HR8gknlZz4W",
-        "outputId": "606c26ad-6a7d-42e0-a014-fc84bc898ef7"
+        "id": "4HR8gknlZz4W"
       },
       "outputs": [],
       "source": [
-        "model\n"
+        "model"
       ]
     },
     {
@@ -298,16 +275,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 38
-        },
-        "id": "IQKOohdpztS7",
-        "outputId": "b0538aca-3916-4d2b-f5f2-f394cdba3dad"
+        "id": "IQKOohdpztS7"
       },
       "outputs": [],
       "source": [
-        "model.sympy()\n"
+        "model.sympy()"
       ]
     },
     {
@@ -323,16 +295,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 38
-        },
-        "id": "GRcxq-TTlpRX",
-        "outputId": "5e13599e-d469-4110-94a4-689023b40717"
+        "id": "GRcxq-TTlpRX"
       },
       "outputs": [],
       "source": [
-        "model.sympy(2)\n"
+        "model.sympy(2)"
       ]
     },
     {
@@ -357,16 +324,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 35
-        },
-        "id": "HFGaNL6tbDgi",
-        "outputId": "260b0db4-862f-4101-c494-8b03756ed126"
+        "id": "HFGaNL6tbDgi"
       },
       "outputs": [],
       "source": [
-        "model.latex()\n"
+        "model.latex()"
       ]
     },
     {
@@ -384,11 +346,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "Vbz4IMsk2NYH",
-        "outputId": "c0eb8aeb-6656-40a2-eeaf-cd733a2593b8"
+        "id": "Vbz4IMsk2NYH"
       },
       "outputs": [],
       "source": [
@@ -396,7 +354,7 @@
         "ypredict_simpler = model.predict(X, 2)\n",
         "\n",
         "print(\"Default selection MSE:\", np.power(ypredict - y, 2).mean())\n",
-        "print(\"Manual selection MSE for index 2:\", np.power(ypredict_simpler - y, 2).mean())\n"
+        "print(\"Manual selection MSE for index 2:\", np.power(ypredict_simpler - y, 2).mean())"
       ]
     },
     {
@@ -430,7 +388,7 @@
       },
       "outputs": [],
       "source": [
-        "y = X[:, 0] ** 4 - 2\n"
+        "y = X[:, 0] ** 4 - 2"
       ]
     },
     {
@@ -448,12 +406,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "PoEkpvYuGUdy",
-        "outputId": "18493db1-67e3-4493-f5e7-19277dd003d9",
-        "scrolled": true
+        "id": "PoEkpvYuGUdy"
       },
       "outputs": [],
       "source": [
@@ -464,23 +417,18 @@
         "    unary_operators=[\"cos\", \"exp\", \"sin\", \"quart(x) = x^4\"],\n",
         "    extra_sympy_mappings={\"quart\": lambda x: x**4},\n",
         ")\n",
-        "model.fit(X, y)\n"
+        "model.fit(X, y)"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 38
-        },
-        "id": "emn2IajKbDgy",
-        "outputId": "7bfade39-b95e-4314-8d46-854e2421c496"
+        "id": "emn2IajKbDgy"
       },
       "outputs": [],
       "source": [
-        "model.sympy()\n"
+        "model.sympy()"
       ]
     },
     {
@@ -582,7 +530,7 @@
         "X = 2 * np.random.rand(N, 5)\n",
         "sigma = np.random.rand(N) * (5 - 0.1) + 0.1\n",
         "eps = sigma * np.random.randn(N)\n",
-        "y = 5 * np.cos(3.5 * X[:, 0]) - 1.3 + eps\n"
+        "y = 5 * np.cos(3.5 * X[:, 0]) - 1.3 + eps"
       ]
     },
     {
@@ -598,18 +546,13 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 298
-        },
-        "id": "sqMqb4nJ5ZR5",
-        "outputId": "7f06a215-f3b6-4053-fc98-5da227a388a3"
+        "id": "sqMqb4nJ5ZR5"
       },
       "outputs": [],
       "source": [
         "plt.scatter(X[:, 0], y, alpha=0.2)\n",
         "plt.xlabel(\"$x_0$\")\n",
-        "plt.ylabel(\"$y$\")\n"
+        "plt.ylabel(\"$y$\")"
       ]
     },
     {
@@ -629,22 +572,18 @@
       },
       "outputs": [],
       "source": [
-        "weights = 1 / sigma ** 2\n"
+        "weights = 1 / sigma ** 2"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "v8WBYtcZbDhC",
-        "outputId": "de926074-1742-4fa3-fe44-188307e9214c"
+        "id": "v8WBYtcZbDhC"
       },
       "outputs": [],
       "source": [
-        "weights[:5]\n"
+        "weights[:5]"
       ]
     },
     {
@@ -660,12 +599,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "a07K3KUjOxcp",
-        "outputId": "629898f5-36aa-4616-bdb7-b41e07619a02",
-        "scrolled": true
+        "id": "a07K3KUjOxcp"
       },
       "outputs": [],
       "source": [
@@ -676,7 +610,7 @@
         "    binary_operators=[\"plus\", \"mult\"],\n",
         "    unary_operators=[\"cos\"],\n",
         ")\n",
-        "model.fit(X, y, weights=weights)\n"
+        "model.fit(X, y, weights=weights)"
       ]
     },
     {
@@ -692,15 +626,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "oHyUbcg6ggmx",
-        "outputId": "a261b520-5a89-42c9-982d-b05b905885fb"
+        "id": "oHyUbcg6ggmx"
       },
       "outputs": [],
       "source": [
-        "model\n"
+        "model"
       ]
     },
     {
@@ -716,19 +646,14 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 38
-        },
-        "id": "PB67POLr8b_L",
-        "outputId": "69e4d94b-a00a-4e59-a781-38571b5f42e0"
+        "id": "PB67POLr8b_L"
       },
       "outputs": [],
       "source": [
         "best_idx = model.equations_.query(\n",
         "    f\"loss < {2 * model.equations_.loss.min()}\"\n",
         ").score.idxmax()\n",
-        "model.sympy(best_idx)\n"
+        "model.sympy(best_idx)"
       ]
     },
     {
@@ -753,18 +678,300 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 265
-        },
-        "id": "ezCC0IkS8zFf",
-        "outputId": "d928f975-3843-4430-93fa-4ae3a18abb51"
+        "id": "ezCC0IkS8zFf"
       },
       "outputs": [],
       "source": [
         "plt.scatter(X[:, 0], y, alpha=0.1)\n",
         "y_prediction = model.predict(X, index=best_idx)\n",
-        "plt.scatter(X[:, 0], y_prediction)\n"
+        "plt.scatter(X[:, 0], y_prediction)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# Multiple outputs"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "For multiple outputs, multiple equations are returned:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "X = 2 * np.random.randn(100, 5)\n",
+        "y = 1 / X[:, [0, 1, 2]]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "model = PySRRegressor(\n",
+        "    binary_operators=[\"+\", \"*\"],\n",
+        "    unary_operators=[\"inv(x) = 1/x\"],\n",
+        "    extra_sympy_mappings={\"inv\": lambda x: 1/x},\n",
+        ")\n",
+        "model.fit(X, y)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "model"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# Julia packages and types"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "PySR uses [SymbolicRegression.jl](https://github.com/MilesCranmer/SymbolicRegression.jl)\n",
+        "as its search backend. This is a pure Julia package, and so can interface easily with any other\n",
+        "Julia package.\n",
+        "For some tasks, it may be necessary to load such a package.\n",
+        "\n",
+        "For example, let's say we wish to discovery the following relationship:\n",
+        "\n",
+        "$$ y = p_{3x + 1} - 5, $$\n",
+        "\n",
+        "where $p_i$ is the $i$th prime number, and $x$ is the input feature.\n",
+        "\n",
+        "Let's see if we can discover this using\n",
+        "the [Primes.jl](https://github.com/JuliaMath/Primes.jl) package.\n",
+        "\n",
+        "First, let's get the Julia backend\n",
+        "Here, we might choose to manually specify unlimited threads, `-O3`,\n",
+        "and `compile_modules=False`, although this will only propagate if Julia has not yet started:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import pysr\n",
+        "jl = pysr.julia_helpers.init_julia(\n",
+        "    julia_kwargs={\"threads\": \"auto\", \"optimize\": 2, \"compiled_modules\": False}\n",
+        ")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n",
+        "\n",
+        "`jl` stores the Julia runtime.\n",
+        "\n",
+        "Now, let's run some Julia code to add the Primes.jl\n",
+        "package to the PySR environment:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "jl.eval(\"\"\"\n",
+        "import Pkg\n",
+        "Pkg.add(\"Primes\")\n",
+        "\"\"\")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "This imports the Julia package manager, and uses it to install\n",
+        "`Primes.jl`. Now let's import `Primes.jl`:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "jl.eval(\"import Primes\")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n",
+        "Now, we define a custom operator:\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "jl.eval(\"\"\"\n",
+        "function p(i::T) where T\n",
+        "    if 0.5 < i < 1000\n",
+        "        return T(Primes.prime(round(Int, i)))\n",
+        "    else\n",
+        "        return T(NaN)\n",
+        "    end\n",
+        "end\n",
+        "\"\"\")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n",
+        "We have created a function `p`, which takes a number `i` of type `T` (e.g., `T=Float64`).\n",
+        "`p` first checks whether the input is between 0.5 and 1000.\n",
+        "If out-of-bounds, it returns `NaN`.\n",
+        "If in-bounds, it rounds it to the nearest integer, computes the corresponding prime number, and then\n",
+        "converts it to the same type as input.\n",
+        "\n",
+        "The equivalent function in Python would be:\n",
+        "\n",
+        "```python\n",
+        "import sympy\n",
+        "\n",
+        "def p(i):\n",
+        "    if 0.5 < i < 1000:\n",
+        "        return float(sympy.prime(int(round(i))))\n",
+        "    else:\n",
+        "        return float(\"nan\")\n",
+        "```\n",
+        "\n",
+        "(However, note that this version assumes 64-bit float input, rather than any input type `T`)\n",
+        "\n",
+        "Next, let's generate a list of primes for our test dataset.\n",
+        "Since we are using PyJulia, we can just call `p` directly to do this:\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "primes = {i: jl.p(i*1.0) for i in range(1, 999)}"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Next, let's use this list of primes to create a dataset of $x, y$ pairs:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import numpy as np\n",
+        "\n",
+        "X = np.random.randint(0, 100, 100)[:, None]\n",
+        "y = [primes[3*X[i, 0] + 1] - 5 + np.random.randn()*0.001 for i in range(100)]"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Note that we have also added a tiny bit of noise to the dataset.\n",
+        "\n",
+        "Finally, let's create a PySR model, and pass the custom operator. We also need to define the sympy equivalent, which we can leave as a placeholder for now:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from pysr import PySRRegressor\n",
+        "import sympy\n",
+        "\n",
+        "class sympy_p(sympy.Function):\n",
+        "    pass\n",
+        "\n",
+        "model = PySRRegressor(\n",
+        "    binary_operators=[\"+\", \"-\", \"*\", \"/\"],\n",
+        "    unary_operators=[\"p\"],\n",
+        "    niterations=20,\n",
+        "    extra_sympy_mappings={\"p\": sympy_p}\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "ee30bd41",
+      "metadata": {},
+      "source": [
+        "We are all set to go! Let's see if we can find the true relation:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "model.fit(X, y)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "if all works out, you should be able to see the true relation (note that the constant offset might not be exactly 1, since it is allowed to round to the nearest integer).\n",
+        "\n",
+        "You can get the sympy version of the best equation with:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "model.sympy()"
       ]
     },
     {
@@ -777,6 +984,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "3hS2kTAbbDhL"
@@ -786,9 +994,9 @@
         "\n",
         "Let's consider a time series problem:\n",
         "\n",
-        "$$ z = y^2,\\quad y = \\frac{1}{100} \\sum(y_i),\\quad y_i = x_{i0}^2 + 6 \\cos(2*x_{i2})$$\n",
+        "$$ z = y^2,\\quad y = \\frac{1}{10} \\sum(y_i),\\quad y_i = x_{i0}^2 + 6 \\cos(2*x_{i2})$$\n",
         "\n",
-        "Imagine our time series is 100 timesteps. That is very hard for symbolic regression, even if we impose the inductive bias of $$z=f(\\sum g(x_i))$$ - it is the square of the number of possible equations!\n",
+        "Imagine our time series is 10 timesteps. That is very hard for symbolic regression, even if we impose the inductive bias of $$z=f(\\sum g(x_i))$$ - it is the square of the number of possible equations!\n",
         "\n",
         "But, as in our paper, **we can break this problem down into parts with a neural network. Then approximate the neural network with the symbolic regression!**\n",
         "\n",
@@ -799,22 +1007,21 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "SXJGXySlbDhL",
-        "outputId": "ea1af43c-3823-4779-f981-f8ccc431d055"
+        "id": "SXJGXySlbDhL"
       },
       "outputs": [],
       "source": [
-        "###### np.random.seed(0)\n",
+        "import numpy as np\n",
+        "\n",
+        "rstate = np.random.RandomState(0)\n",
+        "\n",
         "N = 100000\n",
-        "Nt = 100\n",
-        "X = 6 * np.random.rand(N, Nt, 5) - 3\n",
+        "Nt = 10\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",
-        "X.shape, y.shape\n"
+        "X.shape, y.shape"
       ]
     },
     {
@@ -843,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,
@@ -851,9 +1069,14 @@
       },
       "outputs": [],
       "source": [
-        "hidden = 128\n",
-        "total_steps = 50000\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",
@@ -907,7 +1130,7 @@
         "            ),\n",
         "            \"interval\": \"step\",\n",
         "        }\n",
-        "        return [optimizer], [scheduler]\n"
+        "        return [optimizer], [scheduler]"
       ]
     },
     {
@@ -936,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)\n"
+        "test = DataLoader(test_set, batch_size=256, num_workers=cpu_count(), pin_memory=True)"
       ]
     },
     {
@@ -967,18 +1191,14 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "1ldN0999bDhU",
-        "outputId": "269efa86-3331-4ba8-a763-4f26823c8532"
+        "id": "1ldN0999bDhU"
       },
       "outputs": [],
       "source": [
         "pl.seed_everything(0)\n",
         "model = SumNet()\n",
         "model.total_steps = total_steps\n",
-        "model.max_lr = 1e-2\n"
+        "model.max_lr = 1e-2"
       ]
     },
     {
@@ -994,15 +1214,13 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "33R2nrv-b62w",
-        "outputId": "20a8f626-bb4b-4d08-d8ed-e14fa80b2402"
+        "id": "33R2nrv-b62w"
       },
       "outputs": [],
       "source": [
-        "trainer = pl.Trainer(max_steps=total_steps, gpus=1, benchmark=True)\n"
+        "trainer = pl.Trainer(\n",
+        "    max_steps=total_steps, accelerator=\"gpu\", devices=1\n",
+        ")"
       ]
     },
     {
@@ -1018,40 +1236,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 240,
-          "referenced_widgets": [
-            "8f7ca3dc340c4b6ca7607fd5c94c79d0",
-            "cbf3650d65dc4b8986a569d9700207a2",
-            "b8f6f05741f94fd49901887de806eb1d",
-            "ad40f001306a44ecb54f62ab0adda91f",
-            "e4b89b77f1c94abdbffc3ae9a931a148",
-            "ca1e9af3973845c1b2daee8df04f5050",
-            "d638dbb7e0a846ebaca5487f0f384b75",
-            "441f56016ea143d98f2185b42009cc68",
-            "957b8217f89d449ea226c8d211ca055d",
-            "c5956a7501e649319193f6899e6f94af",
-            "e1cc5d5e17ed43ebbd1420cbfbd06758",
-            "0882e10f5ceb4917b6455b74cf4facfb",
-            "5867c0acc2a84f7196aa782f0dc6d4bc",
-            "a243c86070ac4590a99ec844c2cbf677",
-            "397d0a9190f042d8969d96b766e93b90",
-            "76abe6940c3642ea90bbab0409d47f80",
-            "c3d9456987764530a49f80fd08a8a058",
-            "4419e5228ebb46578d550c3f24096c92",
-            "c1b5d69fe13445179345852b4c5c4a3f",
-            "ff544275991b474981f0e55a01a4739a",
-            "d8747da8b9984e12bf4d7f352a3c0a21",
-            "6c06f7dfa9d84d62b218d8182c164bf9"
-          ]
-        },
-        "id": "TXZdF8k1bDhY",
-        "outputId": "6b1b7f68-5dd8-4613-95a7-776fc71c841f"
+        "id": "TXZdF8k1bDhY"
       },
       "outputs": [],
       "source": [
-        "trainer.fit(model, train_dataloaders=train, val_dataloaders=test)\n"
+        "trainer.fit(model, train_dataloaders=train, val_dataloaders=test)"
       ]
     },
     {
@@ -1069,11 +1258,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "s2sQLla5bDhb",
-        "outputId": "5b48a0b0-6c5e-4e9a-8bfe-e9f888af4b9d"
+        "id": "s2sQLla5bDhb"
       },
       "outputs": [],
       "source": [
@@ -1085,7 +1270,7 @@
         "y_for_pysr = torch.sum(y_i_for_pysr, dim=1) / y_i_for_pysr.shape[1]\n",
         "z_for_pysr = zt[idx]  # Use true values.\n",
         "\n",
-        "X_for_pysr.shape, y_i_for_pysr.shape\n"
+        "X_for_pysr.shape, y_i_for_pysr.shape"
       ]
     },
     {
@@ -1096,33 +1281,90 @@
       "source": [
         "## Learning over the network:\n",
         "\n",
-        "Now, let's fit `g` using PySR:"
+        "Now, let's fit `g` using PySR.\n",
+        "\n",
+        "> **Warning**\n",
+        ">\n",
+        "> First, let's save the data, because sometimes PyTorch and PyJulia's C bindings interfere and cause the colab kernel to crash. If we need to restart, we can just load the data without having to retrain the network:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "nnet_recordings = {\n",
+        "    \"g_input\": X_for_pysr.detach().cpu().numpy().reshape(-1, 5),\n",
+        "    \"g_output\": y_i_for_pysr.detach().cpu().numpy().reshape(-1),\n",
+        "    \"f_input\": y_for_pysr.detach().cpu().numpy().reshape(-1, 1),\n",
+        "    \"f_output\": z_for_pysr.detach().cpu().numpy().reshape(-1),\n",
+        "}\n",
+        "\n",
+        "# Save the data for later use:\n",
+        "import pickle as pkl\n",
+        "\n",
+        "with open(\"nnet_recordings.pkl\", \"wb\") as f:\n",
+        "    pkl.dump(nnet_recordings, f)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "We can now load the data, including after a crash (be sure to re-run the import cells at the top of this notebook, including the one that starts PyJulia)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import pickle as pkl\n",
+        "\n",
+        "nnet_recordings = pkl.load(open(\"nnet_recordings.pkl\", \"rb\"))\n",
+        "f_input = nnet_recordings[\"f_input\"]\n",
+        "f_output = nnet_recordings[\"f_output\"]\n",
+        "g_input = nnet_recordings[\"g_input\"]\n",
+        "g_output = nnet_recordings[\"g_output\"]"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "And now fit using a subsample of the data (symbolic regression only needs a small sample to find the best equation):"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "51QdHVSkbDhc",
-        "outputId": "3058d58a-dbbd-4e78-c810-439a5fca76e6",
-        "scrolled": true
+        "id": "51QdHVSkbDhc"
       },
       "outputs": [],
       "source": [
-        "np.random.seed(1)\n",
-        "tmpX = X_for_pysr.detach().numpy().reshape(-1, 5)\n",
-        "tmpy = y_i_for_pysr.detach().numpy().reshape(-1)\n",
-        "idx2 = np.random.randint(0, tmpy.shape[0], size=3000)\n",
+        "rstate = np.random.RandomState(0)\n",
+        "f_sample_idx = rstate.choice(f_input.shape[0], size=500, replace=False)\n",
         "\n",
         "model = PySRRegressor(\n",
         "    niterations=20,\n",
         "    binary_operators=[\"plus\", \"sub\", \"mult\"],\n",
         "    unary_operators=[\"cos\", \"square\", \"neg\"],\n",
         ")\n",
-        "model.fit(X=tmpX[idx2], y=tmpy[idx2])\n"
+        "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."
       ]
     },
     {
@@ -1135,14 +1377,18 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "6WuaeqyqbDhe"
       },
       "source": [
-        "Recall we are searching for $y_i$ above:\n",
+        "Recall we are searching for $f$ and $g$ such that:\n",
+        "$$z=f(\\sum g(x_i))$$ \n",
+        "which approximates the true relation:\n",
+        "$$ z = y^2,\\quad y = \\frac{1}{10} \\sum(y_i),\\quad y_i = x_{i0}^2 + 6 \\cos(2 x_{i2})$$\n",
         "\n",
-        "$$ z = y^2,\\quad y = \\frac{1}{100} \\sum(y_i),\\quad y_i = x_{i0}^2 + 6 \\cos(2 x_{i2})$$"
+        "Let's see how well we did in recovering $g$:"
       ]
     },
     {
@@ -1153,7 +1399,7 @@
       },
       "outputs": [],
       "source": [
-        "model\n"
+        "model.equations_[[\"complexity\", \"loss\", \"equation\"]]"
       ]
     },
     {
@@ -1162,9 +1408,11 @@
         "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",
-        "Then, we can learn another analytic equation for $z$."
+        "Then, we can learn another analytic equation for $f$."
       ]
     },
     {
@@ -1204,714 +1452,14 @@
   "metadata": {
     "accelerator": "GPU",
     "colab": {
-      "collapsed_sections": [],
       "name": "pysr_demo.ipynb",
       "provenance": []
     },
+    "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.8.9"
-    },
-    "widgets": {
-      "application/vnd.jupyter.widget-state+json": {
-        "0882e10f5ceb4917b6455b74cf4facfb": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "HBoxModel",
-          "state": {
-            "_dom_classes": [],
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "HBoxModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/controls",
-            "_view_module_version": "1.5.0",
-            "_view_name": "HBoxView",
-            "box_style": "",
-            "children": [
-              "IPY_MODEL_5867c0acc2a84f7196aa782f0dc6d4bc",
-              "IPY_MODEL_a243c86070ac4590a99ec844c2cbf677",
-              "IPY_MODEL_397d0a9190f042d8969d96b766e93b90"
-            ],
-            "layout": "IPY_MODEL_76abe6940c3642ea90bbab0409d47f80"
-          }
-        },
-        "397d0a9190f042d8969d96b766e93b90": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "HTMLModel",
-          "state": {
-            "_dom_classes": [],
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "HTMLModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/controls",
-            "_view_module_version": "1.5.0",
-            "_view_name": "HTMLView",
-            "description": "",
-            "description_tooltip": null,
-            "layout": "IPY_MODEL_d8747da8b9984e12bf4d7f352a3c0a21",
-            "placeholder": "​",
-            "style": "IPY_MODEL_6c06f7dfa9d84d62b218d8182c164bf9",
-            "value": " 2120/2442 [00:18&lt;00:02, 116.15it/s, loss=6.39, v_num=14]"
-          }
-        },
-        "4419e5228ebb46578d550c3f24096c92": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "DescriptionStyleModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "DescriptionStyleModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "StyleView",
-            "description_width": ""
-          }
-        },
-        "441f56016ea143d98f2185b42009cc68": {
-          "model_module": "@jupyter-widgets/base",
-          "model_module_version": "1.2.0",
-          "model_name": "LayoutModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "LayoutView",
-            "align_content": null,
-            "align_items": null,
-            "align_self": null,
-            "border": null,
-            "bottom": null,
-            "display": null,
-            "flex": "2",
-            "flex_flow": null,
-            "grid_area": null,
-            "grid_auto_columns": null,
-            "grid_auto_flow": null,
-            "grid_auto_rows": null,
-            "grid_column": null,
-            "grid_gap": null,
-            "grid_row": null,
-            "grid_template_areas": null,
-            "grid_template_columns": null,
-            "grid_template_rows": null,
-            "height": null,
-            "justify_content": null,
-            "justify_items": null,
-            "left": null,
-            "margin": null,
-            "max_height": null,
-            "max_width": null,
-            "min_height": null,
-            "min_width": null,
-            "object_fit": null,
-            "object_position": null,
-            "order": null,
-            "overflow": null,
-            "overflow_x": null,
-            "overflow_y": null,
-            "padding": null,
-            "right": null,
-            "top": null,
-            "visibility": null,
-            "width": null
-          }
-        },
-        "5867c0acc2a84f7196aa782f0dc6d4bc": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "HTMLModel",
-          "state": {
-            "_dom_classes": [],
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "HTMLModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/controls",
-            "_view_module_version": "1.5.0",
-            "_view_name": "HTMLView",
-            "description": "",
-            "description_tooltip": null,
-            "layout": "IPY_MODEL_c3d9456987764530a49f80fd08a8a058",
-            "placeholder": "​",
-            "style": "IPY_MODEL_4419e5228ebb46578d550c3f24096c92",
-            "value": "Epoch 0:  87%"
-          }
-        },
-        "6c06f7dfa9d84d62b218d8182c164bf9": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "DescriptionStyleModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "DescriptionStyleModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "StyleView",
-            "description_width": ""
-          }
-        },
-        "76abe6940c3642ea90bbab0409d47f80": {
-          "model_module": "@jupyter-widgets/base",
-          "model_module_version": "1.2.0",
-          "model_name": "LayoutModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "LayoutView",
-            "align_content": null,
-            "align_items": null,
-            "align_self": null,
-            "border": null,
-            "bottom": null,
-            "display": "inline-flex",
-            "flex": null,
-            "flex_flow": "row wrap",
-            "grid_area": null,
-            "grid_auto_columns": null,
-            "grid_auto_flow": null,
-            "grid_auto_rows": null,
-            "grid_column": null,
-            "grid_gap": null,
-            "grid_row": null,
-            "grid_template_areas": null,
-            "grid_template_columns": null,
-            "grid_template_rows": null,
-            "height": null,
-            "justify_content": null,
-            "justify_items": null,
-            "left": null,
-            "margin": null,
-            "max_height": null,
-            "max_width": null,
-            "min_height": null,
-            "min_width": null,
-            "object_fit": null,
-            "object_position": null,
-            "order": null,
-            "overflow": null,
-            "overflow_x": null,
-            "overflow_y": null,
-            "padding": null,
-            "right": null,
-            "top": null,
-            "visibility": null,
-            "width": "100%"
-          }
-        },
-        "8f7ca3dc340c4b6ca7607fd5c94c79d0": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "HBoxModel",
-          "state": {
-            "_dom_classes": [],
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "HBoxModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/controls",
-            "_view_module_version": "1.5.0",
-            "_view_name": "HBoxView",
-            "box_style": "",
-            "children": [
-              "IPY_MODEL_cbf3650d65dc4b8986a569d9700207a2",
-              "IPY_MODEL_b8f6f05741f94fd49901887de806eb1d",
-              "IPY_MODEL_ad40f001306a44ecb54f62ab0adda91f"
-            ],
-            "layout": "IPY_MODEL_e4b89b77f1c94abdbffc3ae9a931a148"
-          }
-        },
-        "957b8217f89d449ea226c8d211ca055d": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "ProgressStyleModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "ProgressStyleModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "StyleView",
-            "bar_color": null,
-            "description_width": ""
-          }
-        },
-        "a243c86070ac4590a99ec844c2cbf677": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "FloatProgressModel",
-          "state": {
-            "_dom_classes": [],
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "FloatProgressModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/controls",
-            "_view_module_version": "1.5.0",
-            "_view_name": "ProgressView",
-            "bar_style": "",
-            "description": "",
-            "description_tooltip": null,
-            "layout": "IPY_MODEL_c1b5d69fe13445179345852b4c5c4a3f",
-            "max": 2442,
-            "min": 0,
-            "orientation": "horizontal",
-            "style": "IPY_MODEL_ff544275991b474981f0e55a01a4739a",
-            "value": 2120
-          }
-        },
-        "ad40f001306a44ecb54f62ab0adda91f": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "HTMLModel",
-          "state": {
-            "_dom_classes": [],
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "HTMLModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/controls",
-            "_view_module_version": "1.5.0",
-            "_view_name": "HTMLView",
-            "description": "",
-            "description_tooltip": null,
-            "layout": "IPY_MODEL_c5956a7501e649319193f6899e6f94af",
-            "placeholder": "​",
-            "style": "IPY_MODEL_e1cc5d5e17ed43ebbd1420cbfbd06758",
-            "value": " 2/2 [00:00&lt;00:00, 55.34it/s]"
-          }
-        },
-        "b8f6f05741f94fd49901887de806eb1d": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "FloatProgressModel",
-          "state": {
-            "_dom_classes": [],
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "FloatProgressModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/controls",
-            "_view_module_version": "1.5.0",
-            "_view_name": "ProgressView",
-            "bar_style": "",
-            "description": "",
-            "description_tooltip": null,
-            "layout": "IPY_MODEL_441f56016ea143d98f2185b42009cc68",
-            "max": 2,
-            "min": 0,
-            "orientation": "horizontal",
-            "style": "IPY_MODEL_957b8217f89d449ea226c8d211ca055d",
-            "value": 2
-          }
-        },
-        "c1b5d69fe13445179345852b4c5c4a3f": {
-          "model_module": "@jupyter-widgets/base",
-          "model_module_version": "1.2.0",
-          "model_name": "LayoutModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "LayoutView",
-            "align_content": null,
-            "align_items": null,
-            "align_self": null,
-            "border": null,
-            "bottom": null,
-            "display": null,
-            "flex": "2",
-            "flex_flow": null,
-            "grid_area": null,
-            "grid_auto_columns": null,
-            "grid_auto_flow": null,
-            "grid_auto_rows": null,
-            "grid_column": null,
-            "grid_gap": null,
-            "grid_row": null,
-            "grid_template_areas": null,
-            "grid_template_columns": null,
-            "grid_template_rows": null,
-            "height": null,
-            "justify_content": null,
-            "justify_items": null,
-            "left": null,
-            "margin": null,
-            "max_height": null,
-            "max_width": null,
-            "min_height": null,
-            "min_width": null,
-            "object_fit": null,
-            "object_position": null,
-            "order": null,
-            "overflow": null,
-            "overflow_x": null,
-            "overflow_y": null,
-            "padding": null,
-            "right": null,
-            "top": null,
-            "visibility": null,
-            "width": null
-          }
-        },
-        "c3d9456987764530a49f80fd08a8a058": {
-          "model_module": "@jupyter-widgets/base",
-          "model_module_version": "1.2.0",
-          "model_name": "LayoutModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "LayoutView",
-            "align_content": null,
-            "align_items": null,
-            "align_self": null,
-            "border": null,
-            "bottom": null,
-            "display": null,
-            "flex": null,
-            "flex_flow": null,
-            "grid_area": null,
-            "grid_auto_columns": null,
-            "grid_auto_flow": null,
-            "grid_auto_rows": null,
-            "grid_column": null,
-            "grid_gap": null,
-            "grid_row": null,
-            "grid_template_areas": null,
-            "grid_template_columns": null,
-            "grid_template_rows": null,
-            "height": null,
-            "justify_content": null,
-            "justify_items": null,
-            "left": null,
-            "margin": null,
-            "max_height": null,
-            "max_width": null,
-            "min_height": null,
-            "min_width": null,
-            "object_fit": null,
-            "object_position": null,
-            "order": null,
-            "overflow": null,
-            "overflow_x": null,
-            "overflow_y": null,
-            "padding": null,
-            "right": null,
-            "top": null,
-            "visibility": null,
-            "width": null
-          }
-        },
-        "c5956a7501e649319193f6899e6f94af": {
-          "model_module": "@jupyter-widgets/base",
-          "model_module_version": "1.2.0",
-          "model_name": "LayoutModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "LayoutView",
-            "align_content": null,
-            "align_items": null,
-            "align_self": null,
-            "border": null,
-            "bottom": null,
-            "display": null,
-            "flex": null,
-            "flex_flow": null,
-            "grid_area": null,
-            "grid_auto_columns": null,
-            "grid_auto_flow": null,
-            "grid_auto_rows": null,
-            "grid_column": null,
-            "grid_gap": null,
-            "grid_row": null,
-            "grid_template_areas": null,
-            "grid_template_columns": null,
-            "grid_template_rows": null,
-            "height": null,
-            "justify_content": null,
-            "justify_items": null,
-            "left": null,
-            "margin": null,
-            "max_height": null,
-            "max_width": null,
-            "min_height": null,
-            "min_width": null,
-            "object_fit": null,
-            "object_position": null,
-            "order": null,
-            "overflow": null,
-            "overflow_x": null,
-            "overflow_y": null,
-            "padding": null,
-            "right": null,
-            "top": null,
-            "visibility": null,
-            "width": null
-          }
-        },
-        "ca1e9af3973845c1b2daee8df04f5050": {
-          "model_module": "@jupyter-widgets/base",
-          "model_module_version": "1.2.0",
-          "model_name": "LayoutModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "LayoutView",
-            "align_content": null,
-            "align_items": null,
-            "align_self": null,
-            "border": null,
-            "bottom": null,
-            "display": null,
-            "flex": null,
-            "flex_flow": null,
-            "grid_area": null,
-            "grid_auto_columns": null,
-            "grid_auto_flow": null,
-            "grid_auto_rows": null,
-            "grid_column": null,
-            "grid_gap": null,
-            "grid_row": null,
-            "grid_template_areas": null,
-            "grid_template_columns": null,
-            "grid_template_rows": null,
-            "height": null,
-            "justify_content": null,
-            "justify_items": null,
-            "left": null,
-            "margin": null,
-            "max_height": null,
-            "max_width": null,
-            "min_height": null,
-            "min_width": null,
-            "object_fit": null,
-            "object_position": null,
-            "order": null,
-            "overflow": null,
-            "overflow_x": null,
-            "overflow_y": null,
-            "padding": null,
-            "right": null,
-            "top": null,
-            "visibility": null,
-            "width": null
-          }
-        },
-        "cbf3650d65dc4b8986a569d9700207a2": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "HTMLModel",
-          "state": {
-            "_dom_classes": [],
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "HTMLModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/controls",
-            "_view_module_version": "1.5.0",
-            "_view_name": "HTMLView",
-            "description": "",
-            "description_tooltip": null,
-            "layout": "IPY_MODEL_ca1e9af3973845c1b2daee8df04f5050",
-            "placeholder": "​",
-            "style": "IPY_MODEL_d638dbb7e0a846ebaca5487f0f384b75",
-            "value": "Sanity Checking DataLoader 0: 100%"
-          }
-        },
-        "d638dbb7e0a846ebaca5487f0f384b75": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "DescriptionStyleModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "DescriptionStyleModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "StyleView",
-            "description_width": ""
-          }
-        },
-        "d8747da8b9984e12bf4d7f352a3c0a21": {
-          "model_module": "@jupyter-widgets/base",
-          "model_module_version": "1.2.0",
-          "model_name": "LayoutModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "LayoutView",
-            "align_content": null,
-            "align_items": null,
-            "align_self": null,
-            "border": null,
-            "bottom": null,
-            "display": null,
-            "flex": null,
-            "flex_flow": null,
-            "grid_area": null,
-            "grid_auto_columns": null,
-            "grid_auto_flow": null,
-            "grid_auto_rows": null,
-            "grid_column": null,
-            "grid_gap": null,
-            "grid_row": null,
-            "grid_template_areas": null,
-            "grid_template_columns": null,
-            "grid_template_rows": null,
-            "height": null,
-            "justify_content": null,
-            "justify_items": null,
-            "left": null,
-            "margin": null,
-            "max_height": null,
-            "max_width": null,
-            "min_height": null,
-            "min_width": null,
-            "object_fit": null,
-            "object_position": null,
-            "order": null,
-            "overflow": null,
-            "overflow_x": null,
-            "overflow_y": null,
-            "padding": null,
-            "right": null,
-            "top": null,
-            "visibility": null,
-            "width": null
-          }
-        },
-        "e1cc5d5e17ed43ebbd1420cbfbd06758": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "DescriptionStyleModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "DescriptionStyleModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "StyleView",
-            "description_width": ""
-          }
-        },
-        "e4b89b77f1c94abdbffc3ae9a931a148": {
-          "model_module": "@jupyter-widgets/base",
-          "model_module_version": "1.2.0",
-          "model_name": "LayoutModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "LayoutView",
-            "align_content": null,
-            "align_items": null,
-            "align_self": null,
-            "border": null,
-            "bottom": null,
-            "display": "inline-flex",
-            "flex": null,
-            "flex_flow": "row wrap",
-            "grid_area": null,
-            "grid_auto_columns": null,
-            "grid_auto_flow": null,
-            "grid_auto_rows": null,
-            "grid_column": null,
-            "grid_gap": null,
-            "grid_row": null,
-            "grid_template_areas": null,
-            "grid_template_columns": null,
-            "grid_template_rows": null,
-            "height": null,
-            "justify_content": null,
-            "justify_items": null,
-            "left": null,
-            "margin": null,
-            "max_height": null,
-            "max_width": null,
-            "min_height": null,
-            "min_width": null,
-            "object_fit": null,
-            "object_position": null,
-            "order": null,
-            "overflow": null,
-            "overflow_x": null,
-            "overflow_y": null,
-            "padding": null,
-            "right": null,
-            "top": null,
-            "visibility": null,
-            "width": "100%"
-          }
-        },
-        "ff544275991b474981f0e55a01a4739a": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_module_version": "1.5.0",
-          "model_name": "ProgressStyleModel",
-          "state": {
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "ProgressStyleModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "StyleView",
-            "bar_color": null,
-            "description_width": ""
-          }
-        }
-      }
+      "name": "python3"
     }
   },
   "nbformat": 4,
diff --git a/mkdocs.yml b/mkdocs.yml
index 4fbd94c77..d040ef144 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -34,6 +34,7 @@ nav:
     - api.md
     - api-advanced.md
   - backend.md
+  - interactive-docs.md
 
 extra:
   homepage: https://astroautomata.com/PySR
diff --git a/pysr/julia_helpers.py b/pysr/julia_helpers.py
index 202e4cd11..01d9c3c6a 100644
--- a/pysr/julia_helpers.py
+++ b/pysr/julia_helpers.py
@@ -65,7 +65,7 @@ def _get_io_arg(quiet):
     return io_arg
 
 
-def install(julia_project=None, quiet=False):  # pragma: no cover
+def install(julia_project=None, quiet=False, precompile=None):  # pragma: no cover
     """
     Install PyCall.jl and all required dependencies for SymbolicRegression.jl.
 
@@ -78,17 +78,29 @@ def install(julia_project=None, quiet=False):  # pragma: no cover
     processed_julia_project, is_shared = _process_julia_project(julia_project)
     _set_julia_project_env(processed_julia_project, is_shared)
 
+    if precompile == False:
+        os.environ["JULIA_PKG_PRECOMPILE_AUTO"] = "0"
+
     julia.install(quiet=quiet)
-    Main = init_julia(julia_project, quiet=quiet)
+    Main, init_log = init_julia(julia_project, quiet=quiet, return_aux=True)
     io_arg = _get_io_arg(quiet)
 
+    if precompile is None:
+        precompile = init_log["compiled_modules"]
+
+    if not precompile:
+        Main.eval('ENV["JULIA_PKG_PRECOMPILE_AUTO"] = 0')
+
     if is_shared:
         # Install SymbolicRegression.jl:
         _add_sr_to_julia_project(Main, io_arg)
 
     Main.eval("using Pkg")
     Main.eval(f"Pkg.instantiate({io_arg})")
-    Main.eval(f"Pkg.precompile({io_arg})")
+
+    if precompile:
+        Main.eval(f"Pkg.precompile({io_arg})")
+
     if not quiet:
         warnings.warn(
             "It is recommended to restart Python after installing PySR's dependencies,"
@@ -145,7 +157,7 @@ def _check_for_conflicting_libraries():  # pragma: no cover
         )
 
 
-def init_julia(julia_project=None, quiet=False, julia_kwargs=None):
+def init_julia(julia_project=None, quiet=False, julia_kwargs=None, return_aux=False):
     """Initialize julia binary, turning off compiled modules if needed."""
     global julia_initialized
     global julia_kwargs_at_initialization
@@ -191,6 +203,10 @@ def init_julia(julia_project=None, quiet=False, julia_kwargs=None):
         julia_kwargs = {**julia_kwargs, "compiled_modules": False}
         Julia(**julia_kwargs)
 
+    using_compiled_modules = (not "compiled_modules" in julia_kwargs) or julia_kwargs[
+        "compiled_modules"
+    ]
+
     from julia import Main as _Main
 
     Main = _Main
@@ -230,6 +246,8 @@ def init_julia(julia_project=None, quiet=False, julia_kwargs=None):
         julia_kwargs_at_initialization = julia_kwargs
 
     julia_initialized = True
+    if return_aux:
+        return Main, {"compiled_modules": using_compiled_modules}
     return Main
 
 
diff --git a/pysr/sr.py b/pysr/sr.py
index 2238e6420..760ec1c84 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -1493,7 +1493,7 @@ def _run(self, X, y, mutated_params, weights, seed):
         Main = init_julia(self.julia_project, julia_kwargs=julia_kwargs)
 
         if cluster_manager is not None:
-            cluster_manager = _load_cluster_manager(cluster_manager)
+            cluster_manager = _load_cluster_manager(Main, cluster_manager)
 
         if self.update:
             _, is_shared = _process_julia_project(self.julia_project)
@@ -1573,7 +1573,7 @@ def _run(self, X, y, mutated_params, weights, seed):
             complexity_of_constants=self.complexity_of_constants,
             complexity_of_variables=self.complexity_of_variables,
             nested_constraints=nested_constraints,
-            loss=custom_loss,
+            elementwise_loss=custom_loss,
             maxsize=int(self.maxsize),
             output_file=_escape_filename(self.equation_file_),
             npopulations=int(self.populations),
@@ -2231,7 +2231,7 @@ def latex_table(
             if indices is not None:
                 assert isinstance(indices, list)
                 assert isinstance(indices[0], list)
-                assert isinstance(len(indices), self.nout_)
+                assert len(indices) == self.nout_
 
             generator_fnc = generate_multiple_tables
         else:
diff --git a/pysr/test/test.py b/pysr/test/test.py
index 4d718d909..b19d0ebcf 100644
--- a/pysr/test/test.py
+++ b/pysr/test/test.py
@@ -838,13 +838,13 @@ def test_latex_float_precision(self):
         x = sympy.Symbol("x")
         expr = x * 3232.324857384 - 1.4857485e-10
         self.assertEqual(
-            to_latex(expr, prec=2), "3.2 \cdot 10^{3} x - 1.5 \cdot 10^{-10}"
+            to_latex(expr, prec=2), r"3.2 \cdot 10^{3} x - 1.5 \cdot 10^{-10}"
         )
         self.assertEqual(
-            to_latex(expr, prec=3), "3.23 \cdot 10^{3} x - 1.49 \cdot 10^{-10}"
+            to_latex(expr, prec=3), r"3.23 \cdot 10^{3} x - 1.49 \cdot 10^{-10}"
         )
         self.assertEqual(
-            to_latex(expr, prec=8), "3232.3249 x - 1.4857485 \cdot 10^{-10}"
+            to_latex(expr, prec=8), r"3232.3249 x - 1.4857485 \cdot 10^{-10}"
         )
 
     def test_latex_break_long_equation(self):
diff --git a/pysr/version.py b/pysr/version.py
index 6171f73ef..3c98bbfc8 100644
--- a/pysr/version.py
+++ b/pysr/version.py
@@ -1,2 +1,2 @@
-__version__ = "0.11.12"
-__symbolic_regression_jl_version__ = "0.14.4"
+__version__ = "0.11.15"
+__symbolic_regression_jl_version__ = "0.15.2"