tensorflow 2.16.1 breaks tensorflow-probability with Keras 3.0 API

[matthewfeickert]

👋 Yesterday (2024-03-08) tensorflow v2.16.1 was released on PyPI which introduces breaking changes with Keras 3.0. While the tensorflow v2.16.1 release notes provide instructions on how to force TensorFlow to use the Keras 2.0 API, what is the timeline for tensorflow_probability to adopt the Keras 3.0 API?

Minimal failing example:

$ docker run --rm -ti python:3.11 /bin/bash
root@0046f004ed0d:/# python -m venv venv && . venv/bin/activate
(venv) root@0046f004ed0d:/# python -m pip --quiet install --upgrade pip setuptools wheel
(venv) root@0046f004ed0d:/# python -m pip --quiet install --upgrade tensorflow tensorflow-probability
(venv) root@0046f004ed0d:/# python -m pip list | grep tensor
tensorboard                  2.16.2
tensorboard-data-server      0.7.2
tensorflow                   2.16.1
tensorflow-io-gcs-filesystem 0.36.0
tensorflow-probability       0.23.0
(venv) root@0046f004ed0d:/# python -c 'import tensorflow; import tensorflow_probability'
2024-03-09 07:06:02.555141: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-03-09 07:06:02.556772: I external/local_tsl/tsl/cuda/cudart_stub.cc:32] Could not find cuda drivers on your machine, GPU will not be used.
2024-03-09 07:06:02.559547: I external/local_tsl/tsl/cuda/cudart_stub.cc:32] Could not find cuda drivers on your machine, GPU will not be used.
2024-03-09 07:06:02.590424: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-03-09 07:06:03.189032: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/__init__.py", line 20, in <module>
    from tensorflow_probability import substrates
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/substrates/__init__.py", line 17, in <module>
    from tensorflow_probability.python.internal import all_util
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/__init__.py", line 138, in <module>
    dir(globals()[pkg_name])  # Forces loading the package from its lazy loader.
    ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/internal/lazy_loader.py", line 57, in __dir__
    module = self._load()
             ^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/internal/lazy_loader.py", line 40, in _load
    module = importlib.import_module(self.__name__)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/experimental/__init__.py", line 31, in <module>
    from tensorflow_probability.python.experimental import bayesopt
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/experimental/bayesopt/__init__.py", line 17, in <module>
    from tensorflow_probability.python.experimental.bayesopt import acquisition
  File "/venv/lib/python3.11/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 "/venv/lib/python3.11/site-packages/tensorflow_probability/python/experimental/bayesopt/acquisition/expected_improvement.py", line 19, in <module>
    from tensorflow_probability.python.distributions import normal
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/distributions/__init__.py", line 110, in <module>
    from tensorflow_probability.python.distributions.pixel_cnn import PixelCNN
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/distributions/pixel_cnn.py", line 33, in <module>
    from tensorflow_probability.python.layers import weight_norm
  File "/venv/lib/python3.11/site-packages/tensorflow_probability/python/layers/__init__.py", line 27, in <module>
    from tensorflow_probability.python.layers.distribution_layer import CategoricalMixtureOfOneHotCategorical
  File "/venv/lib/python3.11/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__'
(venv) root@0046f004ed0d:/# 

At the moment, tensorflow-probability is still using the Keras 2.0 API in the codebase:

probability/tensorflow_probability/python/layers/distribution_layer.py

Line 69 in 9a14b9b

tf_keras.__internal__.utils.register_symbolic_tensor_type(dtc._TensorCoercible) # pylint: disable=protected-access

[matthewfeickert]

While a new release of tensorflow-probability will be required to not get broken in the way I showed above, from @jburnim's 988f023 it seems like after this tensorflow-probability release tf-keras will be a required dependency for use with tensorflow as well.

This seems like it should be provided through a tensorflow-probability[tensorflow] extra as communicating dependencies to users through error messages is very annoying from the user side.

edit: If I had taken the time to read 988f023 more carefully before posting this, I would have realized that commit also add this extra as tensorflow-probability[tf]:

probability/setup.py

Lines 110 to 112 in 988f023

	extras_require={ # e.g. `pip install tfp-nightly[jax]`
	'jax': ['jax', 'jaxlib'],
	'tf': [TF_PACKAGE, KERAS_PACKAGE],

probability/setup.py

Lines 51 to 56 in 988f023

	if release:
	TF_PACKAGE = 'tensorflow >= 2.15'
	KERAS_PACKAGE = 'tf-keras >= 2.15'
	else:
	TF_PACKAGE = 'tf-nightly'
	KERAS_PACKAGE = 'tf-keras-nightly'

[jburnim]

Thanks for the report!

The immediate breakage should be fixed by the release today of TFP 0.24.0 -- https://github.com/tensorflow/probability/releases/tag/v0.24.0 . TFP 0.24.0 should work if TF 2.16.1 and TF Keras 2.16 are installed.

We do not currently have any plans to migrate from Keras 2 to Keras 3.

[matthewfeickert]

Thanks @jburnim.

We do not currently have any plans to migrate from Keras 2 to Keras 3.

That's useful to know. It seems that tensorflow-probability is making choices to migrate further from tensorflow (or perhaps the other way around). Are there any plans to just split the library development into one JAX based library and one TensorFlow based library? Or is the idea to just split the support internally and then not duplicate the existing codebase until tensorflow drifts enough to just drop support?

Note to other people that are trying to balance supporting supporting tensorflow-probability with tensorflow in a library, you might want to do something like scikit-hep/pyhf#2452.