message: Updating documentation on the subtle issue for tf.where gradients. The original case seems to work now (because of a quirk of the sqrt_grad op) and didn't have any GradientTape things in there. This makes a bit more clearer where the problem is.
diff: @@ -4648,21 +4648,33 @@ def where_v2(condition, x=None, y=None, name=None):
  dtype=int32)>
  Note that if the gradient of either branch of the tf.where generates
  - a NaN, then the gradient of the entire tf.where will be NaN.
  + a NaN, then the gradient of the entire tf.where will be NaN. This is because
  + the gradient calculation for tf.where combines the two branches, for
  + performance reasons.
  +
  A workaround is to use an inner tf.where to ensure the function has
  no asymptote, and to avoid computing a value whose gradient is NaN by
  replacing dangerous inputs with safe inputs.
  Instead of this,
  - >>> y = tf.constant(-1, dtype=tf.float32)
  - >>> tf.where(y > 0, tf.sqrt(y), y)
  - <tf.Tensor: shape=(), dtype=float32, numpy=-1.0>
  -
  - Use this
  -
  - >>> tf.where(y > 0, tf.sqrt(tf.where(y > 0, y, 1)), y)
  - <tf.Tensor: shape=(), dtype=float32, numpy=-1.0>
  + >>> x = tf.constant(0., dtype=tf.float32)
  + >>> with tf.GradientTape() as tape:
  + ... tape.watch(x)
  + ... y = tf.where(x < 1., 0., 1. / x)
  + >>> print(tape.gradient(y, x))
  + tf.Tensor(nan, shape=(), dtype=float32)
  +
  + Although, the `1. / x` values are never used, its gradient is a NaN when x =
  + 0. Instead, we should guard that with another `tf.where`
  +
  + >>> x = tf.constant(0., dtype=tf.float32)
  + >>> with tf.GradientTape() as tape:
  + ... tape.watch(x)
  + ... safe_x = tf.where(tf.equal(x, 0.), 1., x)
  + ... y = tf.where(x < 1., 0., 1. / safe_x)
  + >>> print(tape.gradient(y, x))
  + tf.Tensor(0.0, shape=(), dtype=float32)
  Args:
  condition: A `tf.Tensor` of type `bool`
  