tensorflow 2.16.1 breaks tensorflow-probability with Keras `3.0` API

[matthewfeickert]

tensorflow 2.16.1 breaks tensorflow-probability with Keras `3.0` API

[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]

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'

[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.