tf.GradientTape() doesn't work on sliced outputs

Question

Here is a piece of code which I tried to run:

import tensorflow as tf

a = tf.constant([[1, 2], [2, 3]], dtype=tf.float32)
b = tf.constant([[1, 2], [2, 3]], dtype=tf.float32)

with tf.GradientTape() as tape1, tf.GradientTape() as tape2:
    tape1.watch(a)
    tape2.watch(a)
    
    c = a * b

grad1 = tape1.gradient(c, a)
grad2 = tape2.gradient(c[:, 0], a)
print(grad1)
print(grad2)

And this is the output:

tf.Tensor(
[[1. 2.]
 [2. 3.]], shape=(2, 2), dtype=float32)
None

As you can observe that tf.GradientTape() is not working with sliced outputs. Is there any way around to this?

score 3 · Answer 1 · answered Aug 03 '20 at 08:38

Yes, everything you do to the tensors needs to happen inside the tape context. You can fix it relatively easily like this:

import tensorflow as tf

a = tf.constant([[1, 2], [2, 3]], dtype=tf.float32)
b = tf.constant([[1, 2], [2, 3]], dtype=tf.float32)

with tf.GradientTape() as tape1, tf.GradientTape() as tape2:
    tape1.watch(a)
    tape2.watch(a)
    
    c = a * b
    c_sliced = c[:, 0]

grad1 = tape1.gradient(c, a)
grad2 = tape2.gradient(c_sliced, a)
print(grad1)
print(grad2)

tf.GradientTape() doesn't work on sliced outputs

1 Answers1