diff --git a/notebooks/hfdemo/patch_tsmixer_getting_started.ipynb b/notebooks/hfdemo/patch_tsmixer_getting_started.ipynb
index 82d2be65..4ec46598 100644
--- a/notebooks/hfdemo/patch_tsmixer_getting_started.ipynb
+++ b/notebooks/hfdemo/patch_tsmixer_getting_started.ipynb
@@ -41,7 +41,6 @@
     "# Third Party\n",
     "from transformers import (\n",
     "    EarlyStoppingCallback,\n",
-    "    PatchTSMixerConfig,\n",
     "    PatchTSMixerForPrediction,\n",
     "    Trainer,\n",
     "    TrainingArguments,\n",
diff --git a/notebooks/hfdemo/patch_tst_getting_started.ipynb b/notebooks/hfdemo/patch_tst_getting_started.ipynb
index 56496e74..1a91c219 100644
--- a/notebooks/hfdemo/patch_tst_getting_started.ipynb
+++ b/notebooks/hfdemo/patch_tst_getting_started.ipynb
@@ -28,7 +28,6 @@
     "# Third Party\n",
     "from transformers import (\n",
     "    EarlyStoppingCallback,\n",
-    "    PatchTSTConfig,\n",
     "    PatchTSTForPrediction,\n",
     "    Trainer,\n",
     "    TrainingArguments,\n",
diff --git a/notebooks/hfdemo/ttm_rolling_prediction_getting_started.ipynb b/notebooks/hfdemo/ttm_rolling_prediction_getting_started.ipynb
index d1c37aa4..7af2bfa6 100644
--- a/notebooks/hfdemo/ttm_rolling_prediction_getting_started.ipynb
+++ b/notebooks/hfdemo/ttm_rolling_prediction_getting_started.ipynb
@@ -32,17 +32,14 @@
     }
    ],
    "source": [
-    "import math\n",
     "import os\n",
     "import tempfile\n",
     "\n",
-    "from torch.optim import AdamW\n",
-    "from torch.optim.lr_scheduler import OneCycleLR\n",
-    "from transformers import EarlyStoppingCallback, Trainer, TrainingArguments, set_seed\n",
+    "from transformers import Trainer, TrainingArguments, set_seed\n",
     "\n",
-    "from tsfm_public import TinyTimeMixerForPrediction, TrackingCallback, count_parameters, load_dataset\n",
-    "from tsfm_public.toolkit.visualization import plot_predictions\n",
-    "from tsfm_public.toolkit import RecursivePredictor,RecursivePredictorConfig"
+    "from tsfm_public import TinyTimeMixerForPrediction, load_dataset\n",
+    "from tsfm_public.toolkit import RecursivePredictor, RecursivePredictorConfig\n",
+    "from tsfm_public.toolkit.visualization import plot_predictions"
    ]
   },
   {
diff --git a/notebooks/thinklab/01_patch_tst_pretrain.ipynb b/notebooks/thinklab/01_patch_tst_pretrain.ipynb
index 374bc6c7..d716fffa 100644
--- a/notebooks/thinklab/01_patch_tst_pretrain.ipynb
+++ b/notebooks/thinklab/01_patch_tst_pretrain.ipynb
@@ -95,8 +95,6 @@
    "source": [
     "import pandas as pd\n",
     "from transformers import (\n",
-    "    PatchTSTConfig,\n",
-    "    PatchTSTForPretraining,\n",
     "    Trainer,\n",
     "    TrainingArguments,\n",
     ")\n",
diff --git a/notebooks/tutorial/ttm_channel_mix_finetuning.ipynb b/notebooks/tutorial/ttm_channel_mix_finetuning.ipynb
index fcf85a9c..961e3997 100644
--- a/notebooks/tutorial/ttm_channel_mix_finetuning.ipynb
+++ b/notebooks/tutorial/ttm_channel_mix_finetuning.ipynb
@@ -89,8 +89,8 @@
     "    count_parameters,\n",
     "    get_datasets,\n",
     ")\n",
-    "from tsfm_public.toolkit.visualization import plot_predictions\n",
-    "from tsfm_public.toolkit.lr_finder import optimal_lr_finder"
+    "from tsfm_public.toolkit.lr_finder import optimal_lr_finder\n",
+    "from tsfm_public.toolkit.visualization import plot_predictions"
    ]
   },
   {
diff --git a/notebooks/tutorial/ttm_with_exog_tutorial.ipynb b/notebooks/tutorial/ttm_with_exog_tutorial.ipynb
index 88ad1dbd..e44baadf 100644
--- a/notebooks/tutorial/ttm_with_exog_tutorial.ipynb
+++ b/notebooks/tutorial/ttm_with_exog_tutorial.ipynb
@@ -79,8 +79,8 @@
     "    count_parameters,\n",
     "    get_datasets,\n",
     ")\n",
-    "from tsfm_public.toolkit.visualization import plot_predictions\n",
-    "from tsfm_public.toolkit.lr_finder import optimal_lr_finder"
+    "from tsfm_public.toolkit.lr_finder import optimal_lr_finder\n",
+    "from tsfm_public.toolkit.visualization import plot_predictions"
    ]
   },
   {
diff --git a/tsfm_public/toolkit/__init__.py b/tsfm_public/toolkit/__init__.py
index dff69312..6889faf3 100644
--- a/tsfm_public/toolkit/__init__.py
+++ b/tsfm_public/toolkit/__init__.py
@@ -4,7 +4,7 @@
 from .callbacks import TrackingCallback
 from .data_handling import load_dataset
 from .dataset import ForecastDFDataset, PretrainDFDataset, RegressionDFDataset
+from .recursive_predictor import RecursivePredictor, RecursivePredictorConfig, RecursivePredictorOutput
 from .time_series_forecasting_pipeline import TimeSeriesForecastingPipeline
 from .time_series_preprocessor import TimeSeriesPreprocessor, get_datasets
 from .util import count_parameters
-from .recursive_predictor import RecursivePredictorOutput, RecursivePredictor, RecursivePredictorConfig