Search code examples
pythontensorflowgoogle-colaboratorytensorflow-probability

tensor is unhashable error when sampling from distribution with tensorflow probability (on colab)

I would like to move some of my code to tensorflow distributions and am working through the tutorials on colab. After requesting tensorflow version 2.x on the runtime, the official sample code no longer works.

You can find the colab notebook here, I added the following code at the very beginning:

try:
  # %tensorflow_version only exists in Colab.
  %tensorflow_version 2.x
except Exception:
  pass

I've only ran the code on colab, but I don't think that it's specific to it. If you want to try to reproduce the problem on your local tensorflow installation, this should be the relevant code:

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions

nd = tfd.MultivariateNormalDiag(loc=[0., 10.], scale_diag=[1., 4.])
nd.sample()

The last line, sampling from nd, creates the following error message:

TypeError: Tensor is unhashable if Tensor equality is enabled. Instead, use tensor.experimental_ref() as the key.

Tensorflow probability is version 0.7.0 and tensorflow is version 2.0.0.

Solution

  • Oh, just after posting I found this issue: https://github.com/tensorflow/probability/issues/540

    The problem is fixed by installing a newer version of tensorflow probability. In colab that can be done by executing:

    !pip install tensorflow-probability==0.8.0rc0

    and then restarting the runtime.

    I also checked the local machine here and it has tensorflow-probability 0.8.0. But if you try to install that in colab, it complains about some dependency mismatch. I guess this whole problem will fix itself soon though, when Google rolls out the up-to-date releases of their libraries in these tutorial notebooks.