diff --git a/discussion/examples/TFP_and_Jax.ipynb b/discussion/examples/TFP_and_Jax.ipynb
index 8f58d26aa6..42dfdca5fe 100644
--- a/discussion/examples/TFP_and_Jax.ipynb
+++ b/discussion/examples/TFP_and_Jax.ipynb
@@ -96,8 +96,8 @@
       "source": [
         "# Importing the TFP with Jax backend\n",
         "!pip3 install -q 'tfp-nightly[jax]' tf-nightly-cpu  # We (currently) still require TF, but TF's smaller CPU build will work.\n",
-        "import tensorflow_probability as tfp\n",
-        "tfp = tfp.experimental.substrates.jax\n",
+        "import tensorflow_probability.substrates.jax as tfp\n",
+        "\n",
         "tf = tfp.tf2jax\n",
         "\n",
         "# Standard TFP Imports\n",