Distribution Output is Incompatible with Keras Dense Layers

[i418c]

In the nightly versions of TF and TFP, distributions passed to a Keras dense layer are not subscriptable, causing a crash. See the gist reproducing this here.

[csuter]

Can you make sure the default versions are uninstalled by pip uninstall'ing them. You can print tf.version (and same for tfp) to check what's being imported

On Mon, Jan 8, 2024 at 18:21 i418c @.***> wrote:

In the nightly versions of TF and TFP, distributions passed to a Keras dense layer are not subscriptable, causing a crash. See the gist reproducing this here https://gist.github.com/i418c/d58f02b7f89adf4511792e12cd5ea65f#file-tfp-distribution-unsubscriptable-ipynb .

— Reply to this email directly, view it on GitHub https://github.com/tensorflow/probability/issues/1780, or unsubscribe https://github.com/notifications/unsubscribe-auth/AABG2GLJ3YJZTUTDRDLIPULYNR5PBAVCNFSM6AAAAABBSGWWPKVHI2DSMVQWIX3LMV43ASLTON2WKOZSGA3TCMZXHA4TIMA . You are receiving this because you are subscribed to this thread.Message ID: @.***>

[i418c]

Added both of those steps to the existing gist. The problems persists.

[csuter]

Thanks just wanted to be sure

On Mon, Jan 8, 2024 at 19:18 i418c @.***> wrote:

Added both of those steps to the existing gist. The problems persists.

— Reply to this email directly, view it on GitHub https://github.com/tensorflow/probability/issues/1780#issuecomment-1882030610, or unsubscribe https://github.com/notifications/unsubscribe-auth/AABG2GLR4NSEQL2SVMMIZBDYNSED7AVCNFSM6AAAAABBSGWWPKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTQOBSGAZTANRRGA . You are receiving this because you commented.Message ID: @.***>

[csuter]

This is beyond my limited keras knowledge. From what I know on the TFP side, it should work. the TensorCoercible mixin just means that the distribution can be passed to TF ops (like tf.add or tf.math.exp) and the distribution will be "cast" to a tf.Tensor by a configurable method (default is sample). Typically, all TF ops call tf.convert_to_tensor on their inputs -- this is the mechanism that "coerces" the distribution to a Tensor. Perhaps keras is not doing this. I am heading out on personal leave shortly and won't have a chance to dig into this. Someone on the keras side might have insight, or maybe @jburnim knows enough to advise.

[csuter]

Oh, but a (partial) workaround is to explicitly call tf.convert_to_tensor between your encoder and decoder calls (ie encoded = tf.convert_to_tensor(encoded)). This gives me a new error about symbolic tensors, that I don't understand.

[jburnim]

TFP is not compatible with Keras 3.

To use Keras and TFP together, you must use Keras 2, which can be installed via the tf-keras or tf-keras-nightly package. And then import tf_keras and use, e.g., tf_keras.layers.Dense and tf_keras.optimizers.Adam instead of keras.layers.Dense and keras.optimizers.Adam. It looks like your example runs with these updates -- https://colab.research.google.com/drive/1EifvJDskVjUOYWzdYXVgHELnTOx75ISe?usp=sharing .

For further info, please see:

• https://keras.io/keras_3/
• https://keras.io/getting_started/#tensorflow--keras-2-backwards-compatibility

[i418c]

I didn't realize that Keras 3 would be imported under tf.keras or that TFP wouldn't be compatible. The links were very helpful. Thanks.