Migrate users of tfp.experimental.substrates.jax to import it as

tensorflow_probability.substrates.jax and to use the JAX specific BUILD target. PiperOrigin-RevId: 613275194
tensorflow · Mar 6, 2024 · 8947615 · 8947615
1 parent 3adb838
commit 8947615
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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",