Update tfp.md

No need to import tensorflow as everything is already running in JAX
blackjax-devs · Sep 20, 2022 · 12badae · 12badae
1 parent 138844e
commit 12badae
Showing 1 changed file with 1 addition and 8 deletions.
diff --git a/examples/tfp.md b/examples/tfp.md
@@ -166,21 +166,14 @@ plt.show()
 
 ## Compare Sampling Time with TFP
 
-```{code-cell} ipython3
-import tensorflow.compat.v2 as tf
-
-tf.enable_v2_behavior()
-```
 
 ```{code-cell} ipython3
 %%time
 
 num_results = 500_000
 num_burnin_steps = 0
 
-# Improve performance by tracing the sampler using `tf.function`
-# and compiling it using XLA.
-@tf.function(autograph=False, experimental_compile=True, experimental_relax_shapes=True)
+@jax.jit
 def do_sampling():
     return tfp.mcmc.sample_chain(
         num_results=num_results,