From 93bdfbfe3d2616957e9a8bc39399de68c9e7f0df Mon Sep 17 00:00:00 2001 From: thomaswc Date: Thu, 7 Mar 2024 18:26:43 -0800 Subject: [PATCH] Import JAX substrate through tfp.substrates.jax instead of deprecated experimental.substrates. Also, have the colab binary depend on both the TensorFlow and JAX substrates since various notebooks use each. PiperOrigin-RevId: 613768759 --- .../jupyter_notebooks/TFP_Release_Notebook_0_11_0.ipynb | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow_probability/examples/jupyter_notebooks/TFP_Release_Notebook_0_11_0.ipynb b/tensorflow_probability/examples/jupyter_notebooks/TFP_Release_Notebook_0_11_0.ipynb index 28c7a447fe..057478279b 100644 --- a/tensorflow_probability/examples/jupyter_notebooks/TFP_Release_Notebook_0_11_0.ipynb +++ b/tensorflow_probability/examples/jupyter_notebooks/TFP_Release_Notebook_0_11_0.ipynb @@ -147,7 +147,7 @@ "config.update('jax_enable_x64', True)\n", "\n", "def demo_jax():\n", - " from tensorflow_probability.python.experimental.substrates import jax as tfp\n", + " import tensorflow_probability.substrates.jax as tfp\n", " tfd = tfp.distributions\n", " tfb = tfp.bijectors\n", "\n", @@ -2123,4 +2123,4 @@ "outputs": [] } ] -} \ No newline at end of file +}