NaN gradients of RelaxedBernoulli samples for p near 0 and 1 in recent TFP versions

I'm encountering some numerical issues when evaluating the gradients of a RelaxedBerrnoulli, with parameter p close to either 0 or 1 (close but not super close). Here is a minimal example:

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors

l = 0.99*tf.ones(shape=[],dtype=tf.float32)

with tf.GradientTape() as tape:
  tape.watch(l)
  p = tfd.RelaxedBernoulli(temperature=0.01, probs=l).sample()
dp_dl = tape.gradient(p,l)
print(p, dp_dl)

which returns NaN for the gradient in TFP v0.13 and TF 2.5. Note that I also see the same thing if I directly try to get the gradient of a Transformed Logistic disribution, so if there is a problem it's probably not at the level of RelaxedBernoulli.

What is interesting is that this used to work perfectly fine in previous TFP version (at least in v0.9). For comparison I made these two notebooks with different TFP versions:

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

Looking at the history of both Logistic and RelaxedBernoulli I really don't see an obvious reason for this change in behavior, so would appreciate any hints/thoughts on what might cause these numerical issues :-)

Start by running the minimal TensorFlow Probability example from the issue with the stated TensorFlow and TFP versions, then compare it with the linked working notebook. Inspect the RelaxedBernoulli and transformed Logistic gradient paths to identify the numerical regression. Done means gradients for probabilities near 0 and 1 are finite and a regression case covers the behavior.