ImportError: No module named tensorflow_probability

[murphyyhuang]

Hi, I run into a problem about some environmental problem ...
I could not run your demo either on Colab on my local pc. All of them reports ModuleNotFoundError with respect to 'tensorflow_probability'.
For Colab of shortest_path.ipynb, change the install setting to be install_graph_nets_library = "yes", it just reports error like:

ImportErrorTraceback (most recent call last)
<ipython-input-5-c414b43b0fe0> in <module>()
      8 import time
      9 
---> 10 from graph_nets import graphs
     11 from graph_nets import utils_np
     12 from graph_nets import utils_tf

/usr/local/lib/python2.7/dist-packages/graph_nets/__init__.py in <module>()
     19 from __future__ import print_function
     20 
---> 21 from graph_nets import blocks
     22 from graph_nets import graphs
     23 from graph_nets import modules

/usr/local/lib/python2.7/dist-packages/graph_nets/blocks.py in <module>()
     36 from graph_nets import graphs
     37 from graph_nets import utils_tf
---> 38 import sonnet as snt
     39 import tensorflow as tf
     40 

/usr/local/lib/python2.7/dist-packages/sonnet/__init__.py in <module>()
     40 import semantic_version
     41 
---> 42 from sonnet.python import custom_getters
     43 from sonnet.python.modules import experimental
     44 from sonnet.python.modules import nets

/usr/local/lib/python2.7/dist-packages/sonnet/python/custom_getters/__init__.py in <module>()
     19 from __future__ import print_function
     20 
---> 21 from sonnet.python.custom_getters import bayes_by_backprop
     22 from sonnet.python.custom_getters.context import Context
     23 from sonnet.python.custom_getters.non_trainable import non_trainable

/usr/local/lib/python2.7/dist-packages/sonnet/python/custom_getters/bayes_by_backprop.py in <module>()
     92 
     93 import tensorflow as tf
---> 94 import tensorflow_probability as tfp
     95 
     96 _DEFAULT_SCALE_TRANSFORM = tf.nn.softplus

ImportError: No module named tensorflow_probability

Similar error is witnessed locally.

Thank you for you help!

[jhamrick]

Thanks for catching this! It looks like the version of sonnet changed right after we tested and released this, and there is apparently an upstream bug in either tensorflow or tensorflow_probability. In the meantime, try installing the previous version of sonnet:

pip install -U dm-sonnet==1.23

And that should work for now.

[murphyyhuang]

That works now. Thanks! : -)

[wajahat57]

That also solves the import error:
ModuleNotFoundError: No module named 'Graph'