How to calculate maximum and mean values of Gradient generated from gradient tape with @tf.functions enabled?

Question

you can get gradients by

grad=tape.gradient(loss, variables) 

But I want two values from this list generated- Maximum gradient and Mean of Gradients. all this while @tf.functions enabled. this is what i have been trying-

max_grad=tf.reduce_max(tf.abs(grad_pde))
mean_grad=tf.reduce_mean(tf.abs(grad_data))
lam=max_grad/mean_grad

but this is giving an error when @tf.functions is enabled. On other hand the computations really become slow with the above disabled.

it gave the following error-

The tensor  cannot be accessed from here, because it was defined in FuncGraph() which is out of scope