我想用 python 编写一个 TensorFlow 操作，但我希望它是可微的（能够计算梯度）。
这个问题询问如何在 python 中编写操作，答案建议使用 py_func （没有梯度）：Tensorflow：用 Python 编写操作
TF 文档描述了如何仅从 C++ 代码开始添加操作：https://www.tensorflow.org/versions/r0.10/how_tos/adding_an_op/index.html
就我而言，我正在进行原型设计，因此我不关心它是否在 GPU 上运行，也不关心它是否可以从 TF python API 以外的任何地方使用。
是的，正如 @Yaroslav 的回答中提到的，这是可能的，关键是他引用的链接：here and here。我想通过举一个具体的例子来详细阐述这个答案。
模运算：让我们在tensorflow中实现逐元素求模运算（它已经存在，但它的梯度尚未定义，但对于本示例，我们将从头开始实现它）。
numpy 函数：第一步是定义我们想要对 numpy 数组执行的操作。逐元素求模运算已经在 numpy 中实现，因此很简单：
import numpy as np
def np_mod(x,y):
return (x % y).astype(np.float32)
原因是.astype(np.float32)是因为默认情况下，tensorflow 采用 float32 类型，如果你给它 float64 （numpy 默认值），它会抱怨。
Gradient Function: Next we need to define the gradient function for our opperation for each input of the opperation as tensorflow function. The function needs to take a very specific form. It need to take the tensorflow representation of the opperation op and the gradient of the output grad and say how to propagate the gradients. In our case, the gradients of the mod opperation are easy, the derivative is 1 with respect to the first argument and
with respect to the second (almost everywhere, and infinite at a finite number of spots, but let's ignore that, see https://math.stackexchange.com/questions/1849280/derivative-of-remainder-function-wrt-denominator for details). So we have
def modgrad(op, grad):
x = op.inputs[0] # the first argument (normally you need those to calculate the gradient, like the gradient of x^2 is 2x. )
y = op.inputs[1] # the second argument
return grad * 1, grad * tf.neg(tf.floordiv(x, y)) #the propagated gradient with respect to the first and second argument respectively
grad 函数需要返回一个 n 元组，其中 n 是操作的参数数量。请注意，我们需要返回输入的张量流函数。
制作带有梯度的 TF 函数：正如上面提到的来源中所解释的，有一个 hack 可以使用以下方法定义函数的梯度tf.RegisterGradient [doc] and tf.Graph.gradient_override_map [doc].
复制代码来自harpone我们可以修改tf.py_func函数使其同时定义渐变：
import tensorflow as tf
def py_func(func, inp, Tout, stateful=True, name=None, grad=None):
# Need to generate a unique name to avoid duplicates:
rnd_name = 'PyFuncGrad' + str(np.random.randint(0, 1E+8))
tf.RegisterGradient(rnd_name)(grad) # see _MySquareGrad for grad example
g = tf.get_default_graph()
with g.gradient_override_map({"PyFunc": rnd_name}):
return tf.py_func(func, inp, Tout, stateful=stateful, name=name)
The stateful选项是告诉tensorflow该函数是否总是为相同的输入提供相同的输出（stateful = False），在这种情况下，tensorflow可以简单地表示张量流图，这是我们的情况，并且在大多数情况下可能都是这种情况。
将它们组合在一起：现在我们已经有了所有的部分，我们可以将它们组合在一起：
from tensorflow.python.framework import ops
def tf_mod(x,y, name=None):
with ops.op_scope([x,y], name, "mod") as name:
z = py_func(np_mod,
[x,y],
[tf.float32],
name=name,
grad=modgrad) # <-- here's the call to the gradient
return z[0]
tf.py_func作用于张量列表（并返回张量列表），这就是为什么我们有[x,y]（并返回z[0]）。
现在我们完成了。我们可以测试它。
Test:
with tf.Session() as sess:
x = tf.constant([0.3,0.7,1.2,1.7])
y = tf.constant([0.2,0.5,1.0,2.9])
z = tf_mod(x,y)
gr = tf.gradients(z, [x,y])
tf.initialize_all_variables().run()
print(x.eval(), y.eval(),z.eval(), gr[0].eval(), gr[1].eval())
[ 0.30000001 0.69999999 1.20000005 1.70000005] [ 0.2 0.5 1. 2.9000001] [ 0.10000001 0.19999999 0.20000005 1.70000005] [ 1. 1. 1.1.] [-1。 -1。 -1。 0.]
Success!