Use torch.autograd.grad for a batch of inputs

hongkai-dai: <p>I would like to use torch.autograd.grad to compute the Jacobian matrix for a batch of input.</p>
<p>Say I have a function <code>y=f(x)</code> where <code>y</code> has the same dimension as <code>x</code>. If I have a batch of input <code>x</code>, denoted as x₁, x₂, …, xₙ, where the batch size is <code>n</code>, I want to compute the quantity</p>
<pre><code class="lang-auto">∑ᵢ trace(∂f(xᵢ)/∂xᵢ)
</code></pre>
<p>Namely for each n x n Jacobian matrix ∂f(xᵢ)/∂xᵢ, I want to compute the trace of this matrix, and then sum up over the data index <code>i</code>.</p>
<p>Currently I use torch.autograd.grad to compute it in a for loop</p>
<pre><code class="lang-auto">import torch


def foo(f: torch.nn.Module, xi: torch.Tensor):
    """
    Compute trace(∂f(xᵢ)/∂xᵢ)
    Here xi is a single input data, not a batch or data.
    """
    xi.requires_grad = True
    x_dim = xi.numel()
    return torch.autograd.grad(f(xi), xi, grad_outputs = (torch.eye(x_dim),), retain_graph=True, create_graph=True, is_grads_batched=True)[0].trace()


if __name__ == "__main__":
    x_batch = torch.tensor([[1, 2.], [3, 4.], [5, 6.]])
    f = torch.nn.Sequential(torch.nn.Linear(2, 3), torch.nn.ReLU(), torch.nn.Linear(3, 2))
    trace_sum = 0
    for i in range(x_batch.shape[0]):
        xi = x_batch[i].clone()
        xi.requires_grad = True
        trace_sum += foo(f, xi)
</code></pre>
<p>Is there a way that I don’t need to loop over the input batch dimension through the for loop? Thanks a lot for your kind help in advance!</p>
<p>P.S I also figured that the desired quantity ∑ᵢ trace(∂f(xᵢ)/∂xᵢ) is just</p>
<pre><code class="lang-auto">∑ᵢ ∑ⱼ∂f(xᵢ)[j]/∂xᵢ[j]
</code></pre>
<p>So should be able to do this</p>
<pre><code class="lang-auto">def foo_batch(f: torch.nn.Module, x):
    """
    Compute ∑ᵢ trace(∂f(xᵢ)/∂xᵢ)
    x is a batch of inputs.
    """
    x.requires_grad = True
    y = f(x)
    y_flat = y.reshape((-1,))
    x_flat = x.reshape((-1,))
    # This is memory inefficient as it requires a huge identity matrix as grad_outputs.
    return torch.autograd.grad(y_flat, x_flat, grad_outputs=(torch.eye(x_flat.numel()),), retain_graph=True, create_graph=True, is_grads_batched=True)[0].trace()

if __name__ == "__main__":
    x_batch = torch.tensor([[1, 2.], [3, 4.], [5, 6.]])
    f = torch.nn.Sequential(torch.nn.Linear(2, 3), torch.nn.ReLU(), torch.nn.Linear(3, 2))
    print(foo_batch(f, x_batch))
</code></pre>
<p>But when I run this code, I got the runtime error</p>
<pre><code class="lang-auto">RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior
</code></pre>
<p>I am not sure what this error means, any help would be appreciated.</p>

KFrank: <p>Hi Hongkai!</p>
<aside class="quote no-group" data-username="hongkai-dai" data-post="1" data-topic="181912" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hongkai-dai/48/61210_2.png" class="avatar"> hongkai-dai:</div>
<blockquote>
<pre><code class="lang-auto">    y = f(x)
    y_flat = y.reshape((-1,))
    x_flat = x.reshape((-1,))
</code></pre>
</blockquote>
</aside>
<p><code>y_flat</code> is a function of <code>x</code> (as is <code>y</code>).  However, <code>y_flat</code> is not a function of<br>
<code>x_flat</code> in the sense of the chain of computations recorded by pytorch’s<br>
computation graph.  (Because <code>x_flat</code> is an invertible function of <code>x</code>, we<br>
can understand <code>y_flat</code> to be implicitly a function of <code>x_flat</code>, but autograd<br>
doesn’t know anything about this.)</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior
</code></pre>
</blockquote>
</aside>
<p>The autograd documentation is rather opaque about this, but this error occurs<br>
because <code>y_flat</code> is not (according to the computation graph) a function of<br>
<code>x_flat</code>.</p>
<p>You may tweak <code>foo_batch()</code> so that <code>y_flat</code> is a function of <code>x_flat</code> (and so<br>
that <code>x_flat</code> becomes a leaf node of the computation graph).</p>
<p>As an aside, it <em>might</em> be faster or more memory efficient to use <a href="https://pytorch.org/docs/2.0/generated/torch.func.jacrev.html#torch.func.jacrev" rel="noopener nofollow ugc">jacrev</a> or<br>
<a href="https://pytorch.org/docs/2.0/generated/torch.autograd.functional.jacobian.html" rel="noopener nofollow ugc">jacobian()</a> to compute your jacobian matrix.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; def foo_batchB (f: torch.nn.Module, x):
...     x_flat = x.reshape((-1,))
...     x_flat.requires_grad = True   # x_flat will now be the leaf of the computation graph
...     y = f (x_flat.reshape (x.shape))
...     y_flat = y.reshape((-1,))
...     # This is memory inefficient as it requires a huge identity matrix as grad_outputs.
...     return torch.autograd.grad(y_flat, x_flat, grad_outputs=(torch.eye(x_flat.numel()),), retain_graph=True, create_graph=True, is_grads_batched=True)[0].trace()
...
&gt;&gt;&gt; f = torch.nn.Sequential (torch.nn.Linear (2, 3), torch.nn.ReLU(), torch.nn.Linear (3, 2))
&gt;&gt;&gt; x_batch = torch.tensor ([[1, 2.], [3, 4.], [5, 6.]])
&gt;&gt;&gt;
&gt;&gt;&gt; foo_batchB (f, x_batch)
tensor(-0.0179, grad_fn=&lt;TraceBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; torch.func.jacrev (f) (x_batch).reshape (x_batch.numel(), x_batch.numel()).trace()
tensor(-0.0179, grad_fn=&lt;TraceBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; torch.autograd.functional.jacobian (f, x_batch, vectorize = True).reshape (x_batch.numel(), x_batch.numel()).trace()
tensor(-0.0179)
</code></pre>
<p>[Edit:  Note that using <code>foo_batchB()</code>, <code>javrev</code>, or <code>jacobian()</code> with <code>x_batch</code><br>
will compute the “block-off-diagonal” elements of the full jacobian matrix (even<br>
though they are all zero).  Your original for-loop version avoids this inefficiency<br>
(although it does compute the full per-batch-element jacobian matrices whose<br>
individual off-diagonal elements get ignored when you compute <code>.trace()</code>).  So<br>
your for-loop version could well be the most efficient, especially as the batch<br>
size becomes larger.  It’s conceivable that you could “vectorize” your for-loop<br>
version using <a href="https://pytorch.org/docs/2.0/generated/torch.func.vmap.html" rel="noopener nofollow ugc">vmap(),</a> but I’ve never used <code>vmap()</code>.]</p>
<p>Best.</p>
<p>K. Frank</p>

hongkai-dai: <p>Thanks a lot <a class="mention" href="/u/kfrank">@KFrank</a> for your kind help! The explanation is super clear, much appreciated!</p>

hongkai-dai: <p>In case it helps future users, when using <code>vmap</code>, the code looks like this</p>
<pre><code class="lang-auto">def foo(f: torch.nn.Module, xi: torch.Tensor):
    x_dim = xi.numel()
    return torch.func.jacrev(f)(xi).reshape(x_dim, x_dim).trace()

batched_foo = torch.func.vmap(foo)
batched_foo(x_batch).sum()
</code></pre>

fpaltri: <p>Hi!<br>
I follow your conversation and I have a similar problem, but why if f goes to one dimensional space raise this error and how I can make a gradient of <code>f</code> w.r.t. its input (maybe considering a batch input of size <code>[n_batch, 2]</code>)?</p>
<pre><code class="lang-auto">RuntimeError: If `is_grads_batched=True`, we interpret the first dimension of each grad_output as the batch dimension. The sizes of the remaining dimensions are expected to match the shape of corresponding output, but a mismatch was detected: grad_output[0] has a shape of torch.Size([]) and output[0] has a shape of torch.Size([3]). If you only want some tensors in `grad_output` to be considered batched, consider using vmap.
</code></pre>
<p>You can easily replicate by defining:</p>
<pre><code class="lang-auto">f = torch.nn.Sequential (torch.nn.Linear (2, 3), torch.nn.ReLU(), torch.nn.Linear (3, 1))
</code></pre>