Persistent

def tf_gradient_tape_persistent(x):
    """
    Simple implementation to understand the functioning of gradient tape for chain rule with persistent set to True
    Inputs:
    x: Tensor value 

    Returns:
    EagerTensor: Derivative of y with respect to input tensor x and derivate of z with respect to x
    """
    with tf.GradientTape(persistent=True) as t:
        t.watch(x)
        y=x*x  ## Defining y(x)=x**2
        z=y*y  ## Defining z(y)=y**2
    # Calculating the derivative y with respect to x
    dy_dx=t.gradient(y,x)
    # Calculating the derivative y with respect to x
    dz_dx = t.gradient(z, x)
    return dy_dx,dz_dx

# Run the function for x=5
tmp_x = tf.constant(3.0)
dy_dx,dz_dx = tf_gradient_tape_persistent(tmp_x)
result_zx = dz_dx.numpy()
result_yx= dy_dx.numpy()