tensorflow / probability

Probabilistic reasoning and statistical analysis in TensorFlow

Home Page:https://www.tensorflow.org/probability/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Keras 3 breaks Tensorflow Probability upon import

AjaniStewart opened this issue · comments

Summary of problem

The recent release of Keras 3 breaks TensorFlow Probability at import.

In a fresh Python 3.10 virtual environment, installation of tensorflow v2.15.0, tensorflow-probability v0.23.0, and keras v3 causes a

AttributeError: module 'keras._tf_keras.keras' has no attribute '__internal__'

Steps to reproduce:

$ conda create -n tfp_test python=3.10 pip
$ conda activate tfp_test
$ pip install tensorflow tensorflow_probability
$ pip install --upgrade keras 
$ python -c 'import tensorflow; import tensorflow_probability'

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/__init__.py", line 20, in <module>
    from tensorflow_probability import substrates
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/substrates/__init__.py", line 17, in <module>
    from tensorflow_probability.python.internal import all_util
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/python/__init__.py", line 138, in <module>
    dir(globals()[pkg_name])  # Forces loading the package from its lazy loader.
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/python/internal/lazy_loader.py", line 57, in __dir__
    module = self._load()
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/python/internal/lazy_loader.py", line 40, in _load
    module = importlib.import_module(self.__name__)
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/python/experimental/__init__.py", line 31, in <module>
    from tensorflow_probability.python.experimental import bayesopt
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/python/experimental/bayesopt/__init__.py", line 17, in <module>
    from tensorflow_probability.python.experimental.bayesopt import acquisition
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/python/experimental/bayesopt/acquisition/__init__.py", line 19, in <module>
    from tensorflow_probability.python.experimental.bayesopt.acquisition.expected_improvement import GaussianProcessExpectedImprovement
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/python/experimental/bayesopt/acquisition/expected_improvement.py", line 19, in <module>
    from tensorflow_probability.python.distributions import normal
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/python/distributions/__init__.py", line 110, in <module>
    from tensorflow_probability.python.distributions.pixel_cnn import PixelCNN
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/python/distributions/pixel_cnn.py", line 33, in <module>
    from tensorflow_probability.python.layers import weight_norm
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/python/layers/__init__.py", line 27, in <module>
    from tensorflow_probability.python.layers.distribution_layer import CategoricalMixtureOfOneHotCategorical
  File "/home/ajani/.conda/envs/tfp_test/lib/python3.10/site-packages/tensorflow_probability/python/layers/distribution_layer.py", line 68, in <module>
    tf.keras.__internal__.utils.register_symbolic_tensor_type(dtc._TensorCoercible)  # pylint: disable=protected-access
AttributeError: module 'keras._tf_keras.keras' has no attribute '__internal__'

TensorFlow Probability 0.23.0 is not compatible with Keras 3.

And it looks like TensorFlow 2.15.0 is also not compatible with Keras 3. When I run those pip commands, upgrading Keras gives the error message:

tensorflow 2.15.0 requires keras<2.16,>=2.15.0, but you have keras 3.0.0 which is incompatible.

It will be possible to import TensorFlow Probability 0.24.0 with TensorFlow 2.16.0 and Keras 3 installed. But please note that TensorFlow Probability 0.24.0 will continue to use Keras 2, which means that tf-keras 2.16.0 will also have to be installed.

commented

@jburnim Do you have a rough timeline for when we can expect the fix to be released? If anyone knows of a work-around I would also be grateful.

We plan to release TensorFlow Probability 0.24.0 a week or so after TensorFlow 2.16.0 is released.

I don't know when the TensorFlow team plans to release TF 2.16.0. (It looks like TF 2.16.0-rc0 was made available last week.)

Until then, you should be able to import tensorflow_probabability with Keras 3 installed if you also install tfp-nightly and tf-nightly (and tf-keras-nightly).