tf gradients of multiple losses

Question

Suppose I have a multiple losses defined as

losses = ... # a tensor with shape: (10,)

Now I want to find the gradient each loss over a weight w

for i in range(len(10)):
    grad[i] = tf.gradients(losses[i], w)

Now, can I do this without the above for loop but directly?

user_2455494 · Answer

you can use the tf.map_fn to do map any arbitrary function across the first dimension of a tensor. So something like this should do the trick:

def get_grads(x):
  return tf.gradients(x, w)

tf.map_fun(get_grads, losses)