TensorFlow Probability 0.24.0 upgrades to TensorFlow version 2.16. From the release notes, a substantial change in distribution had to be made to allow this. The tf pypi package now ships with Keras version 3 which is a total rewrite to support other backends (PyTorch & JAX). This is a breaking change for the tfp.layers submodule. I think this is probably because Keras 3 enforces that layers return tensors and not abstract things like distributions. To mitigate this, tfp now ships a separate tensorflow-probability[tfp] pypi release which retains support for tfp.layers. This relies on the tf_keras package which is Keras 2.
To support tfp v0.24.0, I had to make a number of changes.
require tensorflow-probability[tf] rather than simple tensorflow-probability
require tf_keras`
remove any references to tf.keras from the code base
import tf_keras as tfk wherever Keras is required
remove tensorflow.python.keras.engine.data_adapter from the training loop
TensorFlow Probability 0.24.0 upgrades to TensorFlow version 2.16. From the release notes, a substantial change in distribution had to be made to allow this. The tf pypi package now ships with Keras version 3 which is a total rewrite to support other backends (PyTorch & JAX). This is a breaking change for the
tfp.layerssubmodule. I think this is probably because Keras 3 enforces that layers return tensors and not abstract things like distributions. To mitigate this,tfpnow ships a separatetensorflow-probability[tfp]pypi release which retains support fortfp.layers. This relies on thetf_keraspackage which is Keras 2.To support
tfpv0.24.0, I had to make a number of changes.tensorflow-probability[tf]rather than simpletensorflow-probabilitytf.kerasfrom the code baseimport tf_keras as tfkwherever Keras is requiredtensorflow.python.keras.engine.data_adapterfrom the training loop