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: 618931679 Change-Id: Ic47d7ed13e46336fd4e593725be14b6ebe86b1f7
google-deepmind · Mar 25, 2024 · 440df89 · 440df89
1 parent a04759f
commit 440df89
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/acme/agents/jax/ail/losses.py b/acme/agents/jax/ail/losses.py
@@ -21,10 +21,10 @@
 from acme.jax import networks as networks_lib
 import jax
 import jax.numpy as jnp
-import tensorflow_probability as tfp
+import tensorflow_probability.substrates.jax as tfp
 import tree
 
-tfp = tfp.experimental.substrates.jax
+
 tfd = tfp.distributions
 
 # The loss is a function taking the discriminator, its state, the demo