Tape gradient gives wrong output

Question

I am trying to compute gradient using tape.gradient() but it gives me wrong answer. The error is in the lines u_z=tape.gradient(u,z,unconnected_gradients=tf.UnconnectedGradients.ZERO) and two lines that follow it from below code. The function u is not constant in the variables z,f,t but the output from computing tape.gradient(u,z) or tape.gradient(u,t) gives me a None object. If I pass unconnected_gradients=tf.UnconnectedGradients.ZERO as the argument, then I get 0.0 as the derivative, which does not make sense. So one thing that might have gone wrong is that the network gets disconnected but I cannot understand why this happens and how to fix it. I am using tensorflow 2.6.0 and keras 2.6.0. I provide the code and error message below.

import tensorflow as tf
import numpy as np
from tensorflow import keras
import os
from tqdm import trange
import matplotlib.pyplot as plt
# Switch of unnecessary TF warning messages
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

class Model():
    def __init__(self):
        self.optimizer = keras.optimizers.Adam()
        self.initializer = tf.keras.initializers.HeNormal()
        self.batchSize = 500
        self.number_epochs=5000
        
    def NN(self,num_layers = 3, num_neurons = 30):
        model_ = keras.models.Sequential()
        model_.add(keras.layers.Dense(num_neurons,activation='tanh',input_dim=3,kernel_initializer = self.initializer))
        for layer in range(num_layers-1):
            model_.add(keras.layers.Dense(num_neurons,activation='tanh',kernel_initializer=self.initializer))
        model_.add(keras.layers.Dense(1,kernel_initializer=self.initializer))
        return model_
    
    def solve_pde(self,value_function,X,idx):
        z,f,t = X[:,0:1],X[:,1:2],X[:,2:3]
        with tf.GradientTape(persistent=True) as tape:
            u = value_function(tf.concat([z,f,t],axis=1))
        u_z = tape.gradient(u,z,unconnected_gradients=tf.UnconnectedGradients.ZERO)
        u_zz = tape.gradient(u_z,z,unconnected_gradients=tf.UnconnectedGradients.ZERO)
        u_t = tape.gradient(u,t)
        u_pde = u_t +   u_z  +  u_zz - tf.cast(0.5,dtype=tf.float32) * u
        return u_pde
    
    def loss_function(self,batchSize):
        z = tf.linspace(0.001,0.999, 200)
        f = tf.linspace(0.1,0.2, 20)
        z_tile = tf.tile(tf.expand_dims(z,axis=-1),multiples=[20,1])
        f_tile = tf.reshape(tf.repeat(f,200),[-1,1])
        dt = 0.9
        X=tf.concat((z_tile,f_tile,tf.reshape(tf.repeat(dt,z_tile.shape[0]),[-1,1])),axis=1)
        X_pde = tf.concat((z_tile,f_tile,tf.random.uniform(shape=(z_tile.shape[0],1),minval=0,maxval=dt)),axis=1)
        x_star = tf.concat((z_tile,f_tile,tf.reshape(tf.repeat(0.0,z_tile.shape[0]),[-1,1])),axis=1)
        idx = np.random.choice(X.shape[0],batchSize,replace=True)
        loss_e = self.solve_pde(self.value_function_e,X_pde,idx)
        self.value_updated = self.value_function_e(tf.concat[x_star[:,0:1],x_star[:,1:2],x_star[:,2:3]]).numpy().reshape(self.innerStep.Nz,self.innerStep.Nf).transpose()
        return loss_e
    
    @tf.function 
    def training_step(self):
        with tf.GradientTape(persistent=True) as tape:
            loss_e = self.loss_function(self.batchSize)
        grads_valueE = tape.gradient(loss_e,self.theta_valueFunction_e)
        self.optimizer.apply_gradients(zip(grads_valueE,self.theta_valueFunction_e))
        return loss_e
    
    def train_model(self):
        self.value_function_e = self.NN()

        self.theta_valueFunction_e = self.value_function_e.trainable_variables
        
        self.LVF= []

        for epoch in trange(self.number_epochs):
            print(epoch)
            loss_e = self.training_step()
            self.LVF_list.append(loss_e.numpy())
            
        
if __name__=="__main__":
    ext = Model()
    ext.train_model()    
    

The error message along with full traceback is

Traceback (most recent call last):

  File "<ipython-input-26-f5a127c3c9ae>", line 1, in <module>
    runfile('C:/Users/user/Google Drive/S/Research Project4/trial.py', wdir='C:/Users/user/Google Drive/SFI/Research Project4')

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\spyder_kernels\customize\spydercustomize.py", line 827, in runfile
    execfile(filename, namespace)

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\spyder_kernels\customize\spydercustomize.py", line 110, in execfile
    exec(compile(f.read(), filename, 'exec'), namespace)

  File "C:/Users/user/Google Drive/SFI/Research Project4/trial.py", line 85, in <module>
    ext.train_model()

  File "C:/Users/user/Google Drive/SFI/Research Project4/trial.py", line 79, in train_model
    loss_e = self.training_step()

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\tensorflow\python\eager\def_function.py", line 862, in __call__
    return self._python_function(*args, **kwds)

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\tensorflow\python\eager\function.py", line 3985, in bound_method_wrapper
    return wrapped_fn(weak_instance(), *args, **kwargs)

  File "C:/Users/user/Google Drive/SFI/Research Project4/trial.py", line 65, in training_step
    loss_e = self.loss_function(self.batchSize)

  File "C:/Users/user/Google Drive/SFI/Research Project4/trial.py", line 58, in loss_function
    loss_e = self.solve_pde(self.value_function_e,X_pde,idx)

  File "C:/Users/user/Google Drive/SFI/Research Project4/trial.py", line 34, in solve_pde
    u_pde = u_t +   u_z  +  u_zz - tf.cast(0.5,dtype=tf.float32) * u

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\tensorflow\python\ops\math_ops.py", line 1399, in r_binary_op_wrapper
    y, x = maybe_promote_tensors(y, x)

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\tensorflow\python\ops\math_ops.py", line 1335, in maybe_promote_tensors
    ops.convert_to_tensor(tensor, dtype, name="x"))

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\tensorflow\python\profiler\trace.py", line 163, in wrapped
    return func(*args, **kwargs)

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\tensorflow\python\framework\ops.py", line 1566, in convert_to_tensor
    ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\tensorflow\python\framework\constant_op.py", line 346, in _constant_tensor_conversion_function
    return constant(v, dtype=dtype, name=name)

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\tensorflow\python\framework\constant_op.py", line 272, in constant
    allow_broadcast=True)

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\tensorflow\python\framework\constant_op.py", line 283, in _constant_impl
    return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\tensorflow\python\framework\constant_op.py", line 308, in _constant_eager_impl
    t = convert_to_eager_tensor(value, ctx, dtype)

  File "C:\Users\user\AppData\Local\Continuum\anaconda3\lib\site-packages\tensorflow\python\framework\constant_op.py", line 106, in convert_to_eager_tensor
    return ops.EagerTensor(value, ctx.device_name, dtype)

ValueError: Attempt to convert a value (None) with an unsupported type (<class 'NoneType'>) to a Tensor.

Any help is much appreciated. Thank you.

Answer 1

You have 2 problems in your code which prevents you from getting the result you want.

1. If you want to compute higher-order derivatives you have to create nested GradientTape objects
2. GradientTape automatically track variables in its context, if you want to track tensors (as in your case, you want to track z and t) you have to call tape.watch(<my_tensor>) otherwise you will not have gradients for it.

Fixed code:

def solve_pde(self, value_function, X, idx):
    z, f, t = X[:, 0:1], X[:, 1:2], X[:, 2:3]
    with tf.GradientTape(persistent=True) as tape:
        tape.watch(z)
        with tf.GradientTape(persistent=True) as tape2:
            tape2.watch(z)
            tape2.watch(t)
            u = value_function(tf.concat([z, f, t], axis=1))
        u_z = tape2.gradient(u, z)
    u_zz = tape.gradient(u_z, z)
    u_t = tape2.gradient(u, t)
    u_pde = u_t + u_z + u_zz - tf.cast(0.5, dtype=tf.float32) * u
    return u_pde

More on gradient tape can be found in the official documentation: https://www.tensorflow.org/api_docs/python/tf/GradientTape