Write Custom Python-Based Gradient Function for an Operation? (without C++ Implementation)

Question

I\'m trying to write a custom gradient function for \'my_op\' which for the sake of the example contains just a call to tf.identity() (ideally, it could be any graph).

Answer 1

The following seems work fine. Do you have any reason prefering python_grad_func instead?

@tf.function.Defun(tf.float32, tf.float32)
def bprop(x, dy):
  return tf.sigmoid(x)

@tf.function.Defun(tf.float32, grad_func=bprop)
def fprop(x):
  return x  # identity

a = tf.Variable(tf.constant([-5., 4., -3., 2., 1.], dtype=tf.float32))
grad = tf.gradients(fprop(a), [a])                                         

with tf.Session() as sess:                                                             
  sess.run(tf.initialize_all_variables())
  result = sess.run(grad)                                                        

print(result)                                                                  

Answer 2

Note that python_grad_func needs the same interface as ops.RegisterGradient (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/framework/function.py#L349).

Here is the modified code example:

def my_op_grad(op, grad): ### instead of my_op_grad(x)                                                  
    return tf.sigmoid(op.inputs[0])                                              

@function.Defun(a=tf.float32, python_grad_func=my_op_grad)                       
def my_op(a):                                                                    
    return tf.identity(a)                                                        

def main(unused_argv):                                                           

  a = tf.Variable(tf.constant([-5., 4., -3., 2., 1.], dtype=tf.float32))         
  sess = tf.Session()                                                            
  sess.run(tf.initialize_all_variables())                                        

  a = tf.identity(a) #workaround for bug github.com/tensorflow/tensorflow/issues/3710

  grad = tf.gradients(my_op(a), [a])[0]                                          
  result = sess.run(grad)                                                        

  print(result)                                                                  

  sess.close()     

Output:

[ 0.00669286  0.98201376  0.04742587  0.88079709  0.7310586 ]