TFP Layers support in TF2

[nbro]

I've already come across different problems related to the eager execution of the TensorFlow code that involves this package. For example, #620 (an issue I had opened and that I've not yet solved) and #629. So, does TensorFlow probability officially support TensorFlow 2 and eager execution? If not, it is time to update and upgrade this package.

[brianwa84]

TF2 is certainly supported and encouraged, and we mostly expect TFP to work. TFP uses a lot of small ops and still doesn't internally use tf.function, so I expect some speedups when we are able to then that on widely. Thanks for filing what you have so far. We're not a huge team, and especially the layers library has been underloved lately, so please keep reporting issues as you encounter them.

[nbro]

@brianwa84 Does it also support Keras or, at least, did you have Keras in mind when you developed it? I am particularly interested in using Keras with TensorFlow probability. In fact, the issue I had mentioned above, #620, is about a script that attempts to use Keras and TensorFlow probability. If it officially supports Keras, if you could provide examples of programs (especially, a Bayesian CNN, the example https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/bayesian_neural_network.py, but with Keras) that uses TensorFlow probability and Keras (to build, train and test the models), I would really appreciate.

[dynamicwebpaige]

@nbro TensorFlow Probability supports tf.keras layers and initializers, but I agree that this could be better documented - am not seeing an example tutorial on the TFP website.

Thank you for the suggestion! 🙂

@martinwicke

[nbro]

Meanwhile, I've found another similar issue: #416.