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.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 separatetensorflow-probability[tfp]
pypi release which retains support fortfp.layers
. This relies on thetf_keras
package which is Keras 2.To support
tfp
v0.24.0, I had to make a number of changes.tensorflow-probability[tf]
rather than simpletensorflow-probability
tf.keras
from the code baseimport tf_keras as tfk
wherever Keras is requiredtensorflow.python.keras.engine.data_adapter
from the training loop