0

I have found that if I want to use tf.gradients in tensorflow 2 rather than a gradient tape, I can do this by wrapping the code in a tf.function-decorated function. But somehow, I can't take the gradient with respect to a variable this way:

import tensorflow as tf
a = tf.Variable(initial_value=1.0, dtype=tf.float32)
b = 0.01 * a

@tf.function
def get_grads():
  return tf.gradients(b, a)[0]

print(get_grads())

I would expect to get some sort of tensor as a result, a tensor that ought to evaluate to 0.01. But instead I get None. Note that I am running this on Google Colab, so there should not be any issues with the tensorflow version or installation.

What am I doing wrong?

Charles Staats
  • 203
  • 1
  • 9

1 Answers1

1

The op b = 0.01 * a is out the graph created by the tf.function-decorated function.

you can use :

a = tf.Variable(initial_value=1.0, dtype=tf.float32)
@tf.function
def get_grads():
  
   b = 0.01 * a
   return tf.gradients(b, a)

print(get_grads())
Tou You
  • 1,149
  • 8
  • 7
  • So, `tf.function` always creates a new graph, and ops inside it can only see other ops created by the same graph? That explains a lot, actually. (Like issues I've been having differentiating through subgraphs that are themselves wrapped by `tf.function`.) – Charles Staats Aug 06 '20 at 17:29