I have tried on colab with TF version 2.1.0 , 2.2.0-rc2 and was able to reproduce the issue.Please, find the gist here. Thanks!

This is due to a limitation limitation in how gradients are calculated. Unfortunately, it is unlikely to be fixed in the foreseable future.

You can find more detail here, along with a recipe for how to avoid it: https://stackoverflow.com/questions/33712178/tensorflow-nan-bug/42497444#42497444

In short, if the input to a tf.where contains NaNs, the gradient will always be NaN, regardless whether the input is actually used or not, and the workaround is to prevent the inputs from ever containing NaNs.

Are you satisfied with the resolution of your issue?
Yes
No

Shouldn't this be documented with big warning in tf.where docs in this case?

Indeed it should.

@mdanatg Hello, this is my first time contributing to TensofFlow lib. From the thread I gather you would require the tf.where be updated. If it is so can I work on this?

Hello @0x0badc0de , @mdanatg
Should the updated doc contain a something like a warning? or will a small note at the end, about the input not being Nan will do? Also should the workaround for avoiding it also be added to the doc?

@joemaren @anorak-k

Sorry for the delay. Feel free to send a PR - it's only a matter of adding a paragraph to the docstring.

The text should be more in the lines of a warning. Something like: Important: if any of the inputs contain NaN values, etc.. And yes, it should include the workaround as well, which is something in the lines of: instead of tf.where(x, ops_that_can_nan(z), ...), write tf.where(x, ops_that_can_nan(tf.where(x, z, safe_value)), ...).

@mdanatg I have added the change and raised a PR #38467

@mdanatg Thanks for your reply. However, I would like to mention that this behavior also happens when the generated value in the inactive branch is not finite (i.e. inf or -inf). Here is a minimal reproducible example:

import tensorflow as tf

a = tf.Variable(10.)
with tf.GradientTape() as tape:
  out = tf.where(a < 15., a, tf.math.pow(10.0, tf.math.exp(a)))
  grads = tape.gradient(out, a)

print(grads)
# tf.Tensor(nan, shape=(), dtype=float32)

And also if we reverse the condition such that the branch with infinite value is selected, the gradient would be infinite (which is a bit surprising that it does not generate nan instead, like above):

with tf.GradientTape() as tape:
  out = tf.where(a > 15., a, tf.math.pow(10.0, tf.math.exp(a)))
  grads = tape.gradient(out, a)

print(grads)
# tf.Tensor(inf, shape=(), dtype=float32)

So this behavior happens for both nan and infinite values in inactive branch. I wish it wasn't like this, because it's a bit unreasonable and makes it impossible to use user-defined ops/functions which generate extremely large values for some input values; hence, that inner tf.where workaround may not be practical always (unfortunately, even gradient clipping does not help with this, because clipping a nan value produces nan in TF).

CC: @anorak-k for potential consideration in your PR after @mdanatg confirms this.

@mkaze that's true - nan, inf and any other special FP value will disrupt the gradient calculation.

What happens internally is that the gradients are aggregated in this fashion: 1 * <grad of branch taken> + 0 * <grad of branch not taken>. In the former case, you have 0 * inf = nan. In the latter case, you have 1 * inf = inf. I agree it's very confusing, unfortunately a naive fix would add significant overhead to gradient calculations.

Moreover, the forward calculation doesn't need to result in a nan or inf. You can also get weird results if the gradient alone is nan or inf. For example, the cube root function is defined and well-behaved everywhere, but its derivative at zero is infinite. So this will give you a nan gradient too:

a = tf.Variable(0.0)
with tf.GradientTape() as tape:
  out = tf.where(a < 1, a, tf.pow(a, 1.0/3.0))
  grads = tape.gradient(out, a)
print(grads)

I think the tf.where workaround is useful with infinite values as well, so long as the branch not taken is forced to take a gradient that can be safely multiplied by 0. For your example, it would be something like this:

dummy_safe_value = 0
safe_a = tf.where(a > 15., dummy_safe_value, a)
out = tf.where(a > 15., a, tf.math.pow(10.0, tf.math.exp(safe_a)))

I agree that it sometimes can be impractical to do, but in principle it should always be possible as long as you control the inputs to the sensitive functions - all they have to do is force finite values in all the elements that are dropped.

I want to fix the issue #38349

This is due to a limitation limitation in how gradients are calculated. Unfortunately, it is unlikely to be fixed in the foreseable future.

You can find more detail here, along with a recipe for how to avoid it: https://stackoverflow.com/questions/33712178/tensorflow-nan-bug/42497444#42497444

In short, if the input to a tf.where contains NaNs, the gradient will always be NaN, regardless whether the input is actually used or not, and the workaround is to prevent the inputs from ever containing NaNs.

You can simply have it raise a value error if its getting Nan inputs. Or does it not work like that?