Batch-differentiating with respect to single parameter

evanberkowitz: <p>I have a batch of inputs</p>
<pre><code class="lang-python">import torch
x = torch.arange(10)
</code></pre>
<p>and a complicated routine that depends on some parameter C</p>
<pre><code class="lang-python">C = torch.tensor(-5., requires_grad=True)
</code></pre>
<p>I use that routine to compute some outputs <code>y</code>.  Rather than dump the whole complicated routine, I’ll make up something simple that gives the same issues.</p>
<pre><code class="lang-python">y = C * x
</code></pre>
<p>Now, I want to compute the derivative of y with respect to the parameter dy/dC.  In this case of course I know it’s x, but my ‘real’ routine is more involved and the derivative is nontrivial to compute / program.</p>
<p>I know how to do it elementwise,</p>
<pre><code class="lang-python">torch.tensor(tuple(torch.autograd.grad(yi, C, retain_graph=True)[0] for yi in y))
# tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) # the correct answer!
</code></pre>
<p>But in my real problem this is pretty slow.  Can I eliminate the python-level for loop?</p>
<p>I first tried</p>
<pre><code class="lang-python">torch.autograd.grad(y, C)
torch.autograd.grad(y, C.expand(10))
# both give
# RuntimeError: grad can be implicitly created only for scalar outputs
</code></pre>
<p>and then</p>
<pre><code class="lang-python">torch.autograd.grad(y, C, grad_outputs=torch.ones(10), allow_unused=True)
# (tensor(45.),)
</code></pre>
<p>which is <em>close</em>, since that’s <code>x.sum()</code>; I think it did essentially <code>torch.dot(torch.ones(10), dy/dC == x)</code></p>
<p>Well, that’s almost what I want!  Since <code>torch.dot(torch.ones(10), ...) == torch.matmul(torch.ones(10), ...)</code> can I replace <code>torch.ones --&gt; torch.eye</code> and get dy/dC as desired?</p>
<pre><code class="lang-python">torch.autograd.grad(y, C, grad_outputs=torch.eye(10))
# RuntimeError: Mismatch in shape: grad_output[0] has a shape of torch.Size([10, 10]) and output[0] has a shape of torch.Size([10]).
</code></pre>
<p>No, that doesn’t work.  Perhaps if I expand C?</p>
<pre><code class="lang-python">torch.autograd.grad(y, C.expand(10), grad_outputs=torch.ones(10))
# RuntimeError: One of the differentiated Tensors appears to not have been used in the graph.
</code></pre>
<p>hm, I suppose the computational graph can’t “see through” the fact that every entry of <code>C.expand</code> is <code>C</code>.  That makes sense.</p>
<p>Maybe I need to differentiate elementwise and use vmap?</p>
<pre><code class="lang-auto">import functorch
def ddC(y):
    d = torch.autograd.grad(y, C, is_grads_batched=True)
    # d = torch.autograd.grad(y, C.expand(10), is_grads_batched=True) # gives the same error
    return d[0]

functorch.vmap(ddC)(y)
# ValueError: vmap(vjp, in_dims=0, ...)(&lt;inputs&gt;): Got in_dim=0 for an input 
# but the input is of type &lt;class 'NoneType'&gt;. We cannot vmap over non-Tensor 
# arguments, please use None as the respective in_dim
</code></pre>
<p>I don’t understand this error, and I don’t understand where <code>NoneType</code> could be coming from.</p>
<p>Any advice?  How can I get the same result as the element-wise differentiation without a slow python for loop?  Or is that what I should stick to?</p>
<p>Should I be considering forward-mode autodifferentiation instead?</p>

KFrank: <p>Hi Evan!</p>
<aside class="quote no-group" data-username="evanberkowitz" data-post="1" data-topic="164854" 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/evanberkowitz/48/54487_2.png" class="avatar"> evanberkowitz:</div>
<blockquote>
<p>But in my real problem this is pretty slow.  Can I eliminate the python-level for loop?</p>
</blockquote>
</aside>
<p>You can’t eliminate the loop using backward-mode autograd.  Autograd<br>
will let you compute the derivatives of a single scalar (e.g., a loss) with<br>
respect to a batch of variables (i.e., compute the gradient) in a single<br>
pass, but it won’t compute the derivatives of a batch of results with<br>
respect to a single variable in a single pass.</p>
<aside class="quote no-group">
<blockquote>
<p>Should I be considering forward-mode autodifferentiation instead?</p>
</blockquote>
</aside>
<p>Yes (but with the proviso that forward-mode autograd is still in beta /<br>
experimental).</p>
<p>I have not used forward-mode autograd for anything real so I can’t<br>
speak to its stability (nor performance), but here is an illustration applied<br>
to your toy example:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.12.0
&gt;&gt;&gt;
&gt;&gt;&gt; x = torch.arange (10)
&gt;&gt;&gt; C = torch.tensor (-5., requires_grad=True)
&gt;&gt;&gt; y = C * x
&gt;&gt;&gt;
&gt;&gt;&gt; # compute "batch-derivative" with loop over backward-mode autograd
&gt;&gt;&gt; resultA = torch.tensor (tuple (torch.autograd.grad (yi, C, retain_graph=True)[0] for yi in y))
&gt;&gt;&gt; resultA
tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
&gt;&gt;&gt;
&gt;&gt;&gt; C.requires_grad = False
&gt;&gt;&gt;
&gt;&gt;&gt; # compute "batch-derivative" with one pass of forward-mode autograd
&gt;&gt;&gt; Ct = torch.tensor (1.0)
&gt;&gt;&gt; with torch.autograd.forward_ad.dual_level():
...     C_dual = torch.autograd.forward_ad.make_dual (C, Ct)
...     y_dual = C_dual * x
...     resultB = torch.autograd.forward_ad.unpack_dual (y_dual).tangent
...
&gt;&gt;&gt; resultB
tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
&gt;&gt;&gt; torch.equal (resultB, resultA)
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

evanberkowitz: <p>Hi K. Frank,</p>
<p>Thanks for your response.  I will try forward-mode autodiff, we’ll see what happens.  If I have to write code / do a manual loop, so be it.  What’s peculiar to me is that the computational graph for each entry of y is the same graph; shouldn’t it be re-usable?</p>
<p>Can you explain what the <code>is_grads_batched</code> option does for <code>torch.autograd.grad</code> with an example?  The documentation is pretty opaque as to what the actual use-case is, and can you explain why my <code>vmap(ddC)</code> approach failed?  I thought that was the way to avoid the manual loop over y.</p>

KFrank: <p>Hi Evan!</p>
<aside class="quote no-group" data-username="evanberkowitz" data-post="3" data-topic="164854" 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/evanberkowitz/48/54487_2.png" class="avatar"> evanberkowitz:</div>
<blockquote>
<p>What’s peculiar to me is that the computational graph for each entry of y is the same graph; shouldn’t it be re-usable?</p>
</blockquote>
</aside>
<p>To be clear, you are reusing the computation graph.</p>
<p>You run <code>y = C * x</code> once (the forward pass), which creates the computation<br>
graph just once.  Then when you run<br>
<code>torch.autograd.grad(yi, C, retain_graph=True)</code> in your <code>for yi in y</code><br>
loop, you use the computation graph and retain it so that the same<br>
computation graph can be used again in subsequent iterations of the<br>
loop (without recreating it).  (You are performing multiple backward passes,<br>
but through the same, reused computation graph.)</p>
<aside class="quote no-group">
<blockquote>
<p>Can you explain what the <code>is_grads_batched</code> option does<br>
…<br>
and can you explain why my <code>vmap(ddC)</code> approach failed?</p>
</blockquote>
</aside>
<p>I’ve never used <code>is_grads_batched</code> nor <code>vmap</code> (although the documentation<br>
suggests that <code>is_grads_batched</code> uses <code>vmap</code> under the hood).</p>
<p>Best.</p>
<p>K. Frank</p>

KFrank: <p>Hi Evan!</p>
<p>A quick follow-up:</p>
<aside class="quote no-group" data-username="evanberkowitz" data-post="3" data-topic="164854" 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/evanberkowitz/48/54487_2.png" class="avatar"> evanberkowitz:</div>
<blockquote>
<p>Can you explain what the <code>is_grads_batched</code> option does for <code>torch.autograd.grad</code> with an example?</p>
</blockquote>
</aside>
<p>When using <code>is_grads_batched</code> you have to pass <code>torch.autograd.grad()</code><br>
a <em>batch</em> of <code>grad_outputs</code>, packaged as a tensor.  In your use case<br>
<code>torch.eye()</code> suffices:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.12.0
&gt;&gt;&gt;
&gt;&gt;&gt; x = torch.arange (10)
&gt;&gt;&gt; C = torch.tensor (-5., requires_grad=True)
&gt;&gt;&gt; y = C * x
&gt;&gt;&gt;
&gt;&gt;&gt; # compute "batch-derivative" with loop over backward-mode autograd
&gt;&gt;&gt; resultA = torch.tensor (tuple (torch.autograd.grad (yi, C, retain_graph=True)[0] for yi in y))
&gt;&gt;&gt; resultA
tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
&gt;&gt;&gt;
&gt;&gt;&gt; resultB, = torch.autograd.grad (y, C, torch.eye (10), is_grads_batched = True)   # returns a tuple
&gt;&gt;&gt;
&gt;&gt;&gt; resultB
tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
&gt;&gt;&gt; torch.equal (resultB, resultA)
True
</code></pre>
<p>It’s interesting to note that I ran this example on a version 1.12 without<br>
<code>functorch</code> or (to my knowledge) <code>vmap()</code> installed, so I have no idea<br>
how the batch of gradient computations is being vectorized.  (Maybe<br>
there’s a secret copy of <code>vmap()</code> hiding somewhere, or maybe this falls<br>
back to a simple loop with no speed-up.)</p>
<p>Best.</p>
<p>K. Frank</p>

JosueCom: <p>If I have the following operations</p>
<pre><code class="lang-python">x = th.randn(10, 1)
x.requires_grad_()
C = th.randn(1, 2)
y = C*x # shape (10, 2)
</code></pre>
<p>How can I get the Jacobian of each y batch with respect to their corresponding x? Hence, a batched Jacobian of shape (10, 2) or (10, 2, 1)? You can assume the input is of size 1 per batch</p>
<p>I was looking at this page: <a href="https://pytorch.org/tutorials/prototype/vmap_recipe.html" class="inline-onebox" rel="noopener nofollow ugc">torch.vmap — PyTorch Tutorials 1.12.1+cu102 documentation</a> and what was posted above but I do not fully understand what the role of <code>basis_vectors</code> is and forward-mode AD code structure with tangent and primal.</p>