Verify the tensorflow autogradient and the gradient w.r.t matrix

tf_autograd.py

import numpy as np
import pickle
import matplotlib.pyplot as plt
import scipy.io as sio
import tensorflow as tf

tf.enable_eager_execution()
tmm = tf.matmul

def _test():
    
    A = tf.random_uniform((2, 2))
    zeta = tf.random_normal((2, 2))
    W = tf.random_uniform((2, 3))

    with tf.GradientTape() as t:
        t.watch(zeta)
        z = tf.linalg.norm(tf.matmul(tf.matmul(A, zeta),W))**2
    dz_dx = t.gradient(z, zeta)
    
    print(dz_dx) 
    print(2*tmm(tmm(tmm(tmm(tf.transpose(A), A), zeta), W), tf.transpose(W)))

if __name__ == "__main__":
    _test()