<p>After a bit of digging, it seems that it is not trivial to compute per-example gradients in TensorFlow, because this library performs standard back-propagation to compute the gradients (as do other deep learning libraries like PyTorch, Theano and so on), which never actually computes the per-example gradients, it directly obtains the sum of the per-example gradients. Check out <a href="https://github.com/tensorflow/tensorflow/issues/4897" rel="nofollow noreferrer">this discussion</a> for more details.</p>

<p>However, there are some techniques to work around this issue, at least for some use cases. For example, the paper <a href="https://arxiv.org/abs/1510.01799" rel="nofollow noreferrer">Efficient per-example gradient computation</a> by Ian Goodfellow explains how to efficiently compute per-example vectors containing the sum of squared derivatives.  Here is an excerpt from the paper showing the computation (but I highly encourage you read the paper, it is very short):</p>

<p><a href="https://i.stack.imgur.com/CKipm.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/CKipm.png" alt="enter image description here"></a></p>

<p>This algorithm is O(mnp) instead of O(mnp²), where m is the number of examples, n is the number of layers in the neural net, and p is the number of neurons per layer.  So it is much faster than the naive approach (i.e., performing back-prop once per example), especially when p is large, and even more when using a GPU (which speeds up vectorized approaches by a large factor).</p>
Is it possible to compute per-example gradients efficiently in TensorFlow, in just one graph run?
<p>TD;DR: is there a way to evaluate f'(x1), f'(x2), ..., f'(xn) in just one graph run, in a vectorized form? Where f'(x) is the derivative of f(x).</p>

<p>Something like:</p>

<pre><code>x = tf.placeholder(tf.float32, shape=[100])
f = tf.square(x)
f_grad = tf.multiple_gradients(x) # f_grad contains f'(x[0]), f'(x[1]), ...
</code></pre>

<p>More specifically, I'm trying to implement Black Box Stochastic Variational Inference (BBSVI) manually (I know I could use a library like <a href="http://edwardlib.org/" rel="nofollow noreferrer">Edward</a>, but I'm trying to implement it myself).
At one point, I need to compute the mean of f'(x)g(x) across many different values of x (x1, x2, ..., xn), where f(x) and g(x) are two functions, and f'(x) is the derivative of f(x).</p>

<p>Using TensorFlow's autodiff feature, I can compute f'(x1), f'(x2), ..., f'(xn), by simply calling <code>f_prime.eval(feed_dict={x: xi})</code> once for each value xi in (x1, x2, ..., xn).  This is not efficient at all: I would like to use a vectorized form instead, but I'm not sure how to do this.</p>

<p>Perhaps using <code>tf.stop_gradient()</code> somehow? Or using the <code>grad_ys</code> argument in <code>tf.gradients()</code>?</p>