TensorFlow: Can "None" in gradients be automatically replaced with zeros or used in an optimizer?

Question

I used tf.stop_gradient() to turn off the gradient calculation for some of the weights in my neural network. Unfortunately, tf.GradientTape().gradient() assigns the gradient for those weights as None, which does not work with optimizer.apply_gradients. The workaround is to afterwards assign zeros to those gradients.

Is there a better work around?

Is it possible to have tf.GradientTape().gradient() automatically replace None with zeros? Alternatively, is there a way to get the optimizer to work with None in the gradients list?

Answer 1

If you are turning off the gradients for some weights, GradientTape().gradient() automatically sets the gradient for those weights to be None which are not compatible with an optimizer. However, they can be replaced by zeros by setting unconnected_gradients as such:

tape.gradient(loss, weights, unconnected_gradients=tf.UnconnectedGradients.ZERO)

This is then compatible with the optimizer.