How can I get the symbolic gradient [Tensorflow 2.x]

Question

I want to get the symbolic expression for gradient estimation. When I see the output it's quite difficult to understand what's going on.

import tensorflow as tf
@tf.function
def f_k(input_dat):
    y = tf.matmul(tf.sin(input_dat[0]), input_dat[1])
    grads = tf.gradients([y], input_dat)
    # grads = tape.gradient([y], input_dat)
    tf.print('tf >>', grads)
    print('print >>', grads)
    return y, grads


a = tf.Variable([[1., 3.0], [2., 6.0]])
b = tf.Variable([[1.], [2.]])
input_data = [a, b]
y, z = f_k(input_data)
print(y, z)

Output: inside the function

print >> [<tf.Tensor 'gradients/Sin_grad/mul:0' shape=(2, 2) dtype=float32>, <tf.Tensor 'gradients/MatMul_grad/MatMul_1:0' shape=(2, 1) dtype=float32>]
tf >> [[[0.540302277 -1.979985]
 [-0.416146845 1.92034054]], [[1.75076842]
 [-0.138295487]]

As the output, I want which is shown with print:

[<tf.Tensor 'gradients/Sin_grad/mul:0' shape=(2, 2) dtype=float32>, <tf.Tensor 'gradients/MatMul_grad/MatMul_1:0' shape=(2, 1) dtype=float32>]

However, the function always returns the numerical result. Could someone help me to get this symbolic representation of the gradient?

Answer 1

The symbolic representation you want will only work in graph mode. Outside of graph mode, eager execution is enabled by default. What you can do is create a new function to print the values and wrap it with the @tf.function decorator like you are already doing for f_k:

import tensorflow as tf

@tf.function
def f_k(input_dat):
    y = tf.matmul(tf.sin(input_dat[0]), input_dat[1])
    grads = tf.gradients([y], input_dat)
    # grads = tape.gradient([y], input_dat)
    tf.print('tf >>', grads)
    print('print >>', grads)
    return y, grads

a = tf.Variable([[1., 3.0], [2., 6.0]])
b = tf.Variable([[1.], [2.]])
input_data = [a, b]
y, z = f_k(input_data)

@tf.function
def print_symbolic(y, z):
  print(y,z)
  return y, z
y, z = print_symbolic(y, z)

print >> [<tf.Tensor 'gradients/Sin_grad/mul:0' shape=(2, 2) dtype=float32>, <tf.Tensor 'gradients/MatMul_grad/MatMul_1:0' shape=(2, 1) dtype=float32>]
tf >> [[[0.540302277 -1.979985]
 [-0.416146845 1.92034054]], [[1.75076842]
 [-0.138295487]]]
Tensor("y:0", shape=(2, 1), dtype=float32) [<tf.Tensor 'z:0' shape=(2, 2) dtype=float32>, <tf.Tensor 'z_1:0' shape=(2, 1) dtype=float32>]

You could also just access the tensors of your graph:

graph = f_k.get_concrete_function(input_data).graph
print(*[tensor for op in graph.get_operations() for tensor in op.values()], sep="\n")

Tensor("input_dat:0", shape=(), dtype=resource)
Tensor("input_dat_1:0", shape=(), dtype=resource)
Tensor("Sin/ReadVariableOp:0", shape=(2, 2), dtype=float32)
Tensor("Sin:0", shape=(2, 2), dtype=float32)
Tensor("MatMul/ReadVariableOp:0", shape=(2, 1), dtype=float32)
Tensor("MatMul:0", shape=(2, 1), dtype=float32)
Tensor("gradients/Shape:0", shape=(2,), dtype=int32)
Tensor("gradients/grad_ys_0/Const:0", shape=(), dtype=float32)
Tensor("gradients/grad_ys_0:0", shape=(2, 1), dtype=float32)
Tensor("gradients/MatMul_grad/MatMul:0", shape=(2, 2), dtype=float32)
Tensor("gradients/MatMul_grad/MatMul_1:0", shape=(2, 1), dtype=float32)
Tensor("gradients/Sin_grad/Cos:0", shape=(2, 2), dtype=float32)
Tensor("gradients/Sin_grad/mul:0", shape=(2, 2), dtype=float32)
Tensor("StringFormat:0", shape=(), dtype=string)
Tensor("Identity:0", shape=(2, 1), dtype=float32)
Tensor("Identity_1:0", shape=(2, 2), dtype=float32)
Tensor("Identity_2:0", shape=(2, 1), dtype=float32)

Check the docs for more information.