如何使用TensorFlow custom_gradient定义ReLU？

Question

我正在练习使用TensorFlow的custom_gradient装饰器，并且尝试定义一个简单的ReLU。有人会认为，只要将1和否则将x > 0定义为0就容易了。但是，以下代码不会产生与ReLU相同的梯度：

@tf.custom_gradient
def relu(x):
    def grad(dy):
        return tf.cond(tf.reshape(x, []) > 0,
                       lambda: tf.cast(tf.reshape(1, dy.shape), tf.float32),
                       lambda: tf.cast(tf.reshape(0, dy.shape), tf.float32))
    return tf.nn.relu(x), grad

有人可以向我解释为什么ReLU梯度的标准定义不能产生与以下相同的性能：

@tf.custom_gradient
def relu(x):
    def grad(dy):
        return dy
    return tf.nn.relu(x), grad