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

[sachinprasadhs]

Vector-Quantized Variational Autoencoders is not migrated to Keras 3 yet, if you're interested in contributing to the migration, please feel free to create a PR. Thanks!

[pksX01]

Vector-Quantized Variational Autoencoders is not migrated to Keras 3 yet, if you're interested in contributing to the migration, please feel free to create a PR. Thanks!

Yes, I was working on that only but since it is throwing error for Tensorflow-probability, I raised this issue. I think Tensorflow-probability is not compatible with Keras 3 and it is required for this example. Please let me know if you have any workaround for this.