Python/tensorflowban, hogyan lehet egy hálót deriválni parciálisan?

pl. az f(a,b)=a^2+2ab b irányú gradiense a (3,2) pontban így nézne ki.

import tensorflow as tf

a = tf.placeholder('float')

b = tf.placeholder('float')

f = a*a+2*a*b

gradb = tf.gradients(f, b)

sess = tf.Session()

gb = sess.run(gradb, feed_dict={a:3,b:2})

print(gb)

sess.close()

Vagy ha mindkét gradiens kell, akkor a közepén azt csinálod, hogy

grads = tf.gradients(f,[a,b])

...

ga, gb = sess.run(grads, feed_dict={a:3,b:2})