为什么带有 tf.where 的 tf.gradient 返回 None？

Question

I would like to do change input or/and output according to some condition. But the grad is None.

How can I get the gradient and keep the selection like the code below?

input1, input2 = ..., ...
output1, output2 = model(input1), model(input2)

input = tf.where(tf.less(output1, output2), input1, input2)
output = tf.where(tf.less(output1, output2), output1, output2)

grad, = tf.gradient(output, input)

I print the type of input and output, their dimensions are the same as input1/output1(input2/output2). And if I only compute tf.gradient(output1,input1), it is no problem. What's the difference between them?

Answer 1

You can just do this:

input1, input2 = ..., ...
output1, output2 = model(input1), model(input2)
mask = tf.less(output1, output2)
input = tf.where(mask, input1, input2)
output = tf.where(mask, output1, output2)
grad = tf.add(*tf.gradients(output, [input1, input2]))

tf.gradient will return two tensors with zeros on the places where the corresponding input has not been selected. Hence, grad will hold the correct aggregated gradient for input .

Your original approach does not work because, from the point of view of TensorFlow, there is no dependency between input and output . input is computed from input1 and input2 and output is also computed from input1 and input2 , but there is no path in the graph from input to output , so there is no gradient.

Answer 2

in case someone is still stuck with NaN inputs, tf.where can be entirely replaced with:

tf.minimum(tensor_having_nans, value_that_replaces_nans)

Also, tf.maximum works, and gradients don't truncate. In case of 'inf', only tf.minimum works