Hey, all. I’m in the process of getting familiar with TF’s autograph — the series of posts by @pgaleone were super helpful.
If I want to compute and return gradients inside the graph definition using the tf.function decorator, is using tf.GradientTape a sensible approach for TF 2? This is what I have so far for the simple graph illustrated below:
@tf.function def create_graph(x, y, get_grads): with tf.GradientTape(persistent=True) as tape: c = x + y d = y + 1 e = c * d if get_grads == False: return [e, {}] else: de_da = tape.gradient(e, x) de_db = tape.gradient(e, y) de_dc = tape.gradient(e, c) de_dd = tape.gradient(e, d) de_de = tape.gradient(e, e) return [e,{'d_da':de_da, 'd_db':de_db, 'd_dc':de_dc, 'd_dd':de_dd, 'd_de':de_de}] graph = tf.function(create_graph) 