Integrating Keras into JointDistribution

[Sascha Göbel]

Hi,

Is it possible to pass the output of a neural network built with Keras into a joint distribution built with Tensorflow Probability?

More specifically, suppose the last layer of a Keras model is a "tfp.layers.DistributionLambda", so the model returns a  "tfp.distributions.Distribution" instance. How can this be passed into a "JointDistributionCoroutineAutoBatched", so that we can "yield" it as a parameter? The idea being that we aim to subsequently give this a surrogate distribution for variational inference to jointly estimate the neural network and probabilist model (in which the neural network is embedded).

This should be possible, considering that the TFP homepage states that "TensorFlow Probability ... makes it easy to combine probabilistic models and deep learning".

I also posted this question at https://discuss.tensorflow.org/t/help-integrating-keras-model-in-probabilistic-model/12935 where you can find some basic code examples and what we tried so far.

Appreciate any help on this!

Best,

Sascha