Migrate to Keras 3

[dnerini]

Migrate library to use Keras 3 https://keras.io/keras_3/

[frazane]

Migrating to Keras 3 should be relatively easy if we plan to stick to tensorflow as backend. On the other hand, if we want to make our library backend-agnostic, we would have to make quite some changes to the codebase. I see mainly four options at this point.

1. Stick to tensorflow as backend

Pros: less work, less chance of breaking stuff with the rest of mlpp; we just need to adapt tf.keras parts of code
Cons: it's hard to develop with tensorflow

2. Tensorflow probability with JAX backend (via jax substrate)

Pros: we use JAX as backend, which is faster and easier to use to implement new methods; still requires only few adjustments to the codebase (all parts that are tensorflow [not tensorflow probability] related)
Cons: we would still rely on the same API of tensorflow probability so the development experience doesn't really change that much compared to option 1

3. Multi-backend

Going "multi-backend" for me means we could integrate parts of code written in different frameworks, and change the KERAS_BACKEND environment variable depending on what we want to use. We would need to carefully think about which parts of the codebase need to be purely written with Keras ops, and which parts will be backend specific (we could even have something like models/torch_models.py, models/tfp_models.py, etc.).

Pros: we can pick the best framework to develop specific models; everyone can work with their favourite framework; we can take existing implementations of some models and integrate them easily without translating to another framework
Cons: lots of changes required to the library; adds complexity

4. Backend agnostic

Another option would be to make everything truly backend agnostic by just relying on Keras ops. The problem here is that most statistical tools we find in e.g. tensorflow probability or torch (for instance distributions) are missing and would need to be implemented.

Pros: truly integrates the whole Keras 3 ML infrastructure; consistent API
Cons: a LOT of work to be done, either here or via contribution to Keras' codebase

[MicheleCattaneo]

Migrating to Keras 3 is not currently possible while using TensorFlow Probability, as the two are incompatible as of TFP version 0.24.0: TensorFolow Probability releases
This means that option 1 and 2 are not viable for now.

[MicheleCattaneo]

In this gist I showcase how we could use torch for probabilistic layers (in order not to depend on TFP anymore) and use them with Keras3 thanks to the TorchModuleWrapper class.