diff --git a/notebooks/Making a Custom Statespace Model.ipynb b/notebooks/Making a Custom Statespace Model.ipynb
index 28cefc63..a9937581 100644
--- a/notebooks/Making a Custom Statespace Model.ipynb	
+++ b/notebooks/Making a Custom Statespace Model.ipynb	
@@ -19,13 +19,13 @@
     "\n",
     "numpyro.set_host_device_count(4)\n",
     "\n",
-    "import numpy as np\n",
-    "import matplotlib.pyplot as plt\n",
     "import arviz as az\n",
-    "\n",
-    "from pymc_experimental.statespace.core.statespace import PyMCStateSpace\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import pymc as pm\n",
     "import pytensor.tensor as pt\n",
-    "import pymc as pm"
+    "\n",
+    "from pymc_experimental.statespace.core.statespace import PyMCStateSpace"
    ]
   },
   {
@@ -1092,7 +1092,7 @@
    ],
    "source": [
     "az.plot_posterior(\n",
-    "    idata, var_names=[\"ar_params\", \"sigma_x\"], ref_val=true_ar.tolist() + [true_sigma_x]\n",
+    "    idata, var_names=[\"ar_params\", \"sigma_x\"], ref_val=[*true_ar.tolist(), true_sigma_x]\n",
     ");"
    ]
   },
@@ -1169,13 +1169,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from pymc_experimental.statespace.models.utilities import make_default_coords\n",
     "from pymc_experimental.statespace.utils.constants import (\n",
-    "    ALL_STATE_DIM,\n",
     "    ALL_STATE_AUX_DIM,\n",
-    "    OBS_STATE_DIM,\n",
+    "    ALL_STATE_DIM,\n",
     "    SHOCK_DIM,\n",
     ")\n",
-    "from pymc_experimental.statespace.models.utilities import make_default_coords\n",
     "\n",
     "\n",
     "class AutoRegressiveThree(PyMCStateSpace):\n",
diff --git a/notebooks/SARMA Example.ipynb b/notebooks/SARMA Example.ipynb
index ab85c9dc..4aef7b33 100644
--- a/notebooks/SARMA Example.ipynb	
+++ b/notebooks/SARMA Example.ipynb	
@@ -31,20 +31,19 @@
     "\n",
     "numpyro.set_host_device_count(8)\n",
     "\n",
-    "import pymc as pm\n",
-    "from pytensor import tensor as pt\n",
-    "\n",
     "import arviz as az\n",
-    "import statsmodels.api as sm\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import pandas as pd\n",
-    "from scipy import stats\n",
+    "import pymc as pm\n",
+    "import statsmodels.api as sm\n",
     "\n",
-    "import pymc_experimental.statespace as pmss\n",
+    "from pymc.model.transform.optimization import freeze_dims_and_data\n",
+    "from pytensor import tensor as pt\n",
     "from pytensor.link.jax.dispatch import jax_funcify\n",
     "from pytensor.tensor.nlinalg import KroneckerProduct\n",
-    "from pymc.model.transform.optimization import freeze_dims_and_data\n",
+    "\n",
+    "import pymc_experimental.statespace as pmss\n",
     "\n",
     "\n",
     "@jax_funcify.register(KroneckerProduct)\n",
@@ -2582,8 +2581,8 @@
    "source": [
     "fig, ax = plt.subplots()\n",
     "post = az.extract(post_pred).map(np.exp)\n",
-    "hdi = az.hdi(post_pred.map(np.exp))[f\"predicted_posterior_observed\"]\n",
-    "post[f\"predicted_posterior_observed\"].isel(observed_state=0).mean(dim=\"sample\").plot.line(\n",
+    "hdi = az.hdi(post_pred.map(np.exp))[\"predicted_posterior_observed\"]\n",
+    "post[\"predicted_posterior_observed\"].isel(observed_state=0).mean(dim=\"sample\").plot.line(\n",
     "    x=\"time\", ax=ax, add_legend=False, label=\"Posterior Mean, Predicted\"\n",
     ")\n",
     "ax.fill_between(\n",
diff --git a/notebooks/Structural Timeseries Modeling.ipynb b/notebooks/Structural Timeseries Modeling.ipynb
index 01dc8939..21e71426 100644
--- a/notebooks/Structural Timeseries Modeling.ipynb	
+++ b/notebooks/Structural Timeseries Modeling.ipynb	
@@ -22,17 +22,18 @@
     "import sys\n",
     "\n",
     "sys.path.append(\"..\")\n",
-    "from pymc_experimental.statespace import structural as st\n",
-    "from pymc_experimental.statespace.utils.constants import SHORT_NAME_TO_LONG, MATRIX_NAMES\n",
-    "import matplotlib.pyplot as plt\n",
-    "import pymc as pm\n",
     "import arviz as az\n",
-    "import pytensor\n",
-    "import pytensor.tensor as pt\n",
+    "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import pandas as pd\n",
+    "import pymc as pm\n",
+    "import pytensor.tensor as pt\n",
+    "\n",
     "from patsy import dmatrix\n",
     "\n",
+    "from pymc_experimental.statespace import structural as st\n",
+    "from pymc_experimental.statespace.utils.constants import SHORT_NAME_TO_LONG\n",
+    "\n",
     "plt.rcParams.update(\n",
     "    {\n",
     "        \"figure.figsize\": (14, 4),\n",
@@ -61,15 +62,14 @@
    },
    "outputs": [],
    "source": [
-    "from pymc_experimental.statespace.filters.kalman_filter import StandardFilter\n",
-    "from pymc_experimental.statespace.filters.kalman_smoother import KalmanSmoother\n",
+    "from pymc.pytensorf import compile_pymc, inputvars\n",
+    "\n",
     "from pymc_experimental.statespace.filters.distributions import LinearGaussianStateSpace\n",
-    "from pymc.pytensorf import inputvars, compile_pymc\n",
     "\n",
     "\n",
     "def make_numpy_function(mod):\n",
     "    mod = mod.build(verbose=False)\n",
-    "    data = pt.matrix(\"data\", shape=(None, 1))\n",
+    "    pt.matrix(\"data\", shape=(None, 1))\n",
     "    steps = pt.iscalar(\"steps\")\n",
     "    x0, _, c, d, T, Z, R, H, Q = mod._unpack_statespace_with_placeholders()\n",
     "    sequence_names = [x.name for x in [c, d] if x.ndim == 2]\n",
diff --git a/notebooks/VARMAX Example.ipynb b/notebooks/VARMAX Example.ipynb
index 3c8e1d15..f8b37a48 100644
--- a/notebooks/VARMAX Example.ipynb	
+++ b/notebooks/VARMAX Example.ipynb	
@@ -14,21 +14,21 @@
     "\n",
     "numpyro.set_host_device_count(8)\n",
     "\n",
+    "import sys\n",
+    "\n",
+    "import arviz as az\n",
+    "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
-    "import statsmodels.api as sm\n",
     "import pandas as pd\n",
-    "\n",
     "import pymc as pm\n",
     "import pytensor.tensor as pt\n",
-    "import arviz as az\n",
-    "\n",
-    "import matplotlib.pyplot as plt\n",
-    "import sys\n",
+    "import statsmodels.api as sm\n",
     "\n",
     "sys.path.append(\"..\")\n",
-    "import pymc_experimental.statespace as pmss\n",
     "import re\n",
     "\n",
+    "import pymc_experimental.statespace as pmss\n",
+    "\n",
     "config = {\n",
     "    \"figure.figsize\": [12.0, 4.0],\n",
     "    \"figure.dpi\": 72.0 * 2,\n",
@@ -679,7 +679,7 @@
     "    new_labels = []\n",
     "    for label in axis.yaxis.get_majorticklabels():\n",
     "        old_text = \"[\" + label.get_text().split(\"[\")[-1]\n",
-    "        labels = eval(re.sub(\"([\\d\\w]+)\", '\"\\g<1>\"', old_text))\n",
+    "        labels = eval(re.sub(r\"([\\d\\w]+)\", r'\"\\g<1>\"', old_text))\n",
     "        lag, other_var = labels\n",
     "        new_text = f\"L{lag}.{other_var}\"\n",
     "        new_labels.append(new_text)\n",
diff --git a/notebooks/discrete_markov_chain.ipynb b/notebooks/discrete_markov_chain.ipynb
index 8204b3f1..fd518bfc 100644
--- a/notebooks/discrete_markov_chain.ipynb
+++ b/notebooks/discrete_markov_chain.ipynb
@@ -2,6 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
+   "id": "7fb27b941602401d91542211134fc71a",
    "metadata": {},
    "source": [
     "# Discrete Markov Chain Distribution"
@@ -10,6 +11,7 @@
   {
    "cell_type": "code",
    "execution_count": 1,
+   "id": "acae54e37e7d407bbb7b55eff062a284",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -20,16 +22,17 @@
   {
    "cell_type": "code",
    "execution_count": 2,
+   "id": "9a63283cbaf04dbcab1f6479b197f3a8",
    "metadata": {},
    "outputs": [],
    "source": [
     "import arviz as az\n",
+    "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "import pandas as pd\n",
     "import pymc as pm\n",
     "import pytensor\n",
-    "import pandas as pd\n",
     "import statsmodels.api as sm\n",
-    "import matplotlib.pyplot as plt\n",
     "\n",
     "from matplotlib import ticker as mtick\n",
     "\n",
@@ -38,6 +41,7 @@
   },
   {
    "cell_type": "markdown",
+   "id": "8dd0d8092fe74a7c96281538738b07e2",
    "metadata": {},
    "source": [
     "## Demonstration of API "
@@ -46,6 +50,7 @@
   {
    "cell_type": "code",
    "execution_count": 3,
+   "id": "72eea5119410473aa328ad9291626812",
    "metadata": {},
    "outputs": [
     {
@@ -66,6 +71,7 @@
   {
    "cell_type": "code",
    "execution_count": 4,
+   "id": "8edb47106e1a46a883d545849b8ab81b",
    "metadata": {},
    "outputs": [
     {
@@ -131,6 +137,7 @@
   },
   {
    "cell_type": "markdown",
+   "id": "10185d26023b46108eb7d9f57d49d2b3",
    "metadata": {},
    "source": [
     "Dims of the output are `(batch, time)`."
@@ -139,6 +146,7 @@
   {
    "cell_type": "code",
    "execution_count": 5,
+   "id": "8763a12b2bbd4a93a75aff182afb95dc",
    "metadata": {},
    "outputs": [
     {
@@ -167,6 +175,7 @@
   },
   {
    "cell_type": "markdown",
+   "id": "7623eae2785240b9bd12b16a66d81610",
    "metadata": {},
    "source": [
     "## Parameter Recovery\n",
@@ -177,6 +186,7 @@
   {
    "cell_type": "code",
    "execution_count": 6,
+   "id": "7cdc8c89c7104fffa095e18ddfef8986",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -201,6 +211,7 @@
   {
    "cell_type": "code",
    "execution_count": 7,
+   "id": "b118ea5561624da68c537baed56e602f",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -210,6 +221,7 @@
   {
    "cell_type": "code",
    "execution_count": 8,
+   "id": "938c804e27f84196a10c8828c723f798",
    "metadata": {
     "scrolled": false
    },
@@ -287,6 +299,7 @@
   {
    "cell_type": "code",
    "execution_count": 9,
+   "id": "504fb2a444614c0babb325280ed9130a",
    "metadata": {},
    "outputs": [
     {
@@ -306,6 +319,7 @@
   },
   {
    "cell_type": "markdown",
+   "id": "59bbdb311c014d738909a11f9e486628",
    "metadata": {},
    "source": [
     "## Hidden markov model\n",
@@ -318,6 +332,7 @@
   {
    "cell_type": "code",
    "execution_count": 10,
+   "id": "b43b363d81ae4b689946ece5c682cd59",
    "metadata": {},
    "outputs": [
     {
@@ -348,6 +363,7 @@
   {
    "cell_type": "code",
    "execution_count": 11,
+   "id": "8a65eabff63a45729fe45fb5ade58bdc",
    "metadata": {},
    "outputs": [
     {
@@ -524,6 +540,7 @@
   {
    "cell_type": "code",
    "execution_count": 12,
+   "id": "c3933fab20d04ec698c2621248eb3be0",
    "metadata": {
     "scrolled": false
    },
@@ -594,6 +611,7 @@
   {
    "cell_type": "code",
    "execution_count": 13,
+   "id": "4dd4641cc4064e0191573fe9c69df29b",
    "metadata": {},
    "outputs": [
     {
@@ -758,6 +776,7 @@
   },
   {
    "cell_type": "markdown",
+   "id": "8309879909854d7188b41380fd92a7c3",
    "metadata": {},
    "source": [
     "It's quite difficult to sample this model, so I use a very high `target_accept` and a very large number of draws. It seems to be hard to get a sufficient number of effective samples for rare rates without a lot of draws. "
@@ -766,6 +785,7 @@
   {
    "cell_type": "code",
    "execution_count": 14,
+   "id": "3ed186c9a28b402fb0bc4494df01f08d",
    "metadata": {},
    "outputs": [
     {
@@ -848,6 +868,7 @@
   },
   {
    "cell_type": "markdown",
+   "id": "cb1e1581032b452c9409d6c6813c49d1",
    "metadata": {},
    "source": [
     "### Post-Estimation Diagnostics\n",
@@ -858,6 +879,7 @@
   {
    "cell_type": "code",
    "execution_count": 15,
+   "id": "379cbbc1e968416e875cc15c1202d7eb",
    "metadata": {},
    "outputs": [
     {
@@ -887,6 +909,7 @@
   },
   {
    "cell_type": "markdown",
+   "id": "277c27b1587741f2af2001be3712ef0d",
    "metadata": {},
    "source": [
     "...but the trace plots look great! "
@@ -895,6 +918,7 @@
   {
    "cell_type": "code",
    "execution_count": 16,
+   "id": "db7b79bc585a40fcaf58bf750017e135",
    "metadata": {},
    "outputs": [
     {
@@ -915,6 +939,7 @@
   },
   {
    "cell_type": "markdown",
+   "id": "916684f9a58a4a2aa5f864670399430d",
    "metadata": {},
    "source": [
     "Even after 20,000 draws, we only have about 500 samples for the transition probabilities to the more rare state 1."
@@ -923,6 +948,7 @@
   {
    "cell_type": "code",
    "execution_count": 17,
+   "id": "1671c31a24314836a5b85d7ef7fbf015",
    "metadata": {},
    "outputs": [
     {
@@ -1133,6 +1159,7 @@
   },
   {
    "cell_type": "markdown",
+   "id": "33b0902fd34d4ace834912fa1002cf8e",
    "metadata": {},
    "source": [
     "## Comparison with Statsmodels"
@@ -1141,6 +1168,7 @@
   {
    "cell_type": "code",
    "execution_count": 18,
+   "id": "f6fa52606d8c4a75a9b52967216f8f3f",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1155,6 +1183,7 @@
   {
    "cell_type": "code",
    "execution_count": 19,
+   "id": "f5a1fa73e5044315a093ec459c9be902",
    "metadata": {},
    "outputs": [
     {
@@ -1196,6 +1225,7 @@
   },
   {
    "cell_type": "markdown",
+   "id": "cdf66aed5cc84ca1b48e60bad68798a8",
    "metadata": {},
    "source": [
     "## Posterior Prediction"
@@ -1204,6 +1234,7 @@
   {
    "cell_type": "code",
    "execution_count": 20,
+   "id": "28d3efd5258a48a79c179ea5c6759f01",
    "metadata": {},
    "outputs": [
     {
@@ -1266,6 +1297,7 @@
   {
    "cell_type": "code",
    "execution_count": 21,
+   "id": "3f9bc0b9dd2c44919cc8dcca39b469f8",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1279,6 +1311,7 @@
   {
    "cell_type": "code",
    "execution_count": 22,
+   "id": "0e382214b5f147d187d36a2058b9c724",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1288,6 +1321,7 @@
   {
    "cell_type": "code",
    "execution_count": 23,
+   "id": "5b09d5ef5b5e4bb6ab9b829b10b6a29f",
    "metadata": {},
    "outputs": [
     {
@@ -1317,6 +1351,7 @@
   },
   {
    "cell_type": "markdown",
+   "id": "a50416e276a0479cbe66534ed1713a40",
    "metadata": {},
    "source": [
     "## Regime Inference\n",
@@ -1327,6 +1362,7 @@
   {
    "cell_type": "code",
    "execution_count": 24,
+   "id": "46a27a456b804aa2a380d5edf15a5daf",
    "metadata": {},
    "outputs": [
     {
@@ -1364,6 +1400,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "1944c39560714e6e80c856f20744a8e5",
    "metadata": {},
    "outputs": [],
    "source": []
diff --git a/notebooks/latent_approx.ipynb b/notebooks/latent_approx.ipynb
index ea6fa422..38aa958d 100644
--- a/notebooks/latent_approx.ipynb
+++ b/notebooks/latent_approx.ipynb
@@ -25,13 +25,9 @@
     }
    ],
    "source": [
-    "import arviz as az\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
-    "import pymc as pm\n",
-    "from aesara import tensor as at\n",
-    "\n",
-    "import pymc_experimental as pmx"
+    "import pymc as pm"
    ]
   },
   {
@@ -42,7 +38,7 @@
    "outputs": [],
    "source": [
     "from pymc_experimental.gp import HSGP, KarhunenLoeveExpansion, ProjectedProcess\n",
-    "from pymc_experimental.gp.latent_approx import ExpQuad, Matern12, Matern32, Matern52"
+    "from pymc_experimental.gp.latent_approx import ExpQuad"
    ]
   },
   {
diff --git a/notebooks/marginalized_changepoint_model.ipynb b/notebooks/marginalized_changepoint_model.ipynb
index e0bd9cb1..ed3f5ab4 100644
--- a/notebooks/marginalized_changepoint_model.ipynb
+++ b/notebooks/marginalized_changepoint_model.ipynb
@@ -6,11 +6,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import pymc as pm\n",
-    "from pymc_experimental.marginal_model import MarginalModel\n",
-    "import pandas as pd\n",
+    "import arviz as az\n",
     "import numpy as np\n",
-    "import arviz as az"
+    "import pandas as pd\n",
+    "import pymc as pm\n",
+    "\n",
+    "from pymc_experimental.marginal_model import MarginalModel"
    ]
   },
   {