小编Hem*_*ant的帖子

我试图了解如何使用@tf.custom_gradient TensorFlow 1.7 中提供的函数来提供向量相对于向量的自定义梯度。下面的代码是解决以下问题的最小工作示例dz/dx。

y=Ax
z=||y|| ²

另外，此附图描述了手动计算所期望的解决方案

如果我不使用，@tf.custom_gradient那么 TensorFlow 会按预期给出所需的解决方案。我的问题是如何为 y=Ax 提供自定义渐变？我们知道，dy/dx = A^T如上面的附件所示，它显示了与 TensorFlow 输出相匹配的计算步骤。

import tensorflow as tf

#I want to write custom gradient for this function f1
def f1(A,x):
    y=tf.matmul(A,x,name='y')
    return y

#for y= Ax, the derivative is: dy/dx= transpose(A)
@tf.custom_gradient
def f2(A,x):
    y=f1(A,x)
    def grad(dzByDy): # dz/dy = 2y reaches here correctly.
        dzByDx=tf.matmul(A,dzByDy,transpose_a=True) 
        return dzByDx
    return y,grad


x= tf.constant([[1.],[0.]],name='x')
A= tf.constant([ [1., 2.], [3., 4.]],name='A')

y=f1(A,x) # This …