Tensorflow Probability is breaking with Keras 3

[pksX01]

Issue Type

Bug

Source

source

Keras Version

Keras 3.5.0

Custom Code

No

OS Platform and Distribution

Mac OS Sonoma 14.6.1

Python version

3.10.9

GPU model and memory

No response

Current Behavior?

vq_vae.py is getting failed while running with Keras 3 (in both cases - without any changes and with keras 3 backend agnostic changes) because of tensorflow_probability. It seems tensorflow_probability is not compatible with Keras 3.

AttributeError: module 'tensorflow._api.v2.compat.v2.__internal__' has no attribute 'register_load_context_function'. Did you mean: 'register_call_context_function'?

Standalone code to reproduce the issue or tutorial link

Simply running vq_vae.py with Keras 3 will produce this error.

Relevant log output

No response