New versions of TFP cause NaN gradients for RelaxedBernoulli near p=0 and p=1

[EiffL]

I couldn't reproduce one of the figures of the paper (from this notebook: https://github.com/DifferentiableUniverseInitiative/DHOD/blob/master/nb/Zheng2007_demo.ipynb), specifically the figure comparing gradients.

Turns out, it's due to the RelaxedBernoulli being quite sensitive near 0 and 1 in newer versions of TFP, I'm trying to pin down exactly what changed, but checked that with previous versions (0.9) we can get nice gradients even close to the edges: https://colab.research.google.com/drive/1xkXYtC3ER1z0r25fkUa2pwV_b7OlQBnP?usp=sharing

[EiffL]

I made a more clean-cut example:

• working with TFP=0.9 and TF=2.1
• not working with TFP=0.13 and TF=2.5

[EiffL]

couldn't really find an explanation, so opened an issue upstream tensorflow/probability#1393

[EiffL]

This is partially fixed by #36 which introduces a fix to bypass the problematic sigmoid in tfp, but I'm keeping this issue open to remember to remove that fix when TFP will be fixed