Computing jacobian column-wise with Pytorch, and without loop

ergias: <p>In Pytorch, I have a multi-dimension solution stored on each column of a tensor <code>u</code>, and for the training of the PINN I’m developing, I use this function to compute the derivative column-wise.</p>
<pre><code class="lang-py">def dt(self, u, t):
    # return batch_jacobian(u, t)
    N = u.shape[1]
    u_t = [] # initializing the time derivatives
    # for each columns (i), we compute the time derivative ui_t
    for i in range(N):
        ui_t = torch.autograd.grad(u[:, i], t,
                                   grad_outputs=torch.ones_like(u[:, i]),
                                   retain_graph=True,
                                   create_graph=True,
                                   allow_unused=True)[0]
        u_t.append(ui_t) # time derivatives are stored in u_t
    u_t = torch.cat(tuple(u_t), dim=1) # we concateneate all the derivatives
    return u_t
</code></pre>
<p>but it involves a <code>for</code> loop that I would like to remove.</p>
<p>In my research about it, I found this <a href="https://discuss.pytorch.org/t/computing-batch-jacobian-efficiently/80771">topic</a> with the same issue, but I wasn’t able to fix my code with the proposed solutions.</p>
<p>Here is the definition of the functional that is used for the training of the neural network :</p>
<pre><code class="lang-py">def functional_f(self, t, theta):
    self.A = torch.tensor(theta[:,1:], requires_grad=False).float().to(self.device)
    self.mu = torch.tensor(theta[:,0], requires_grad=False).float().to(self.device)

    u = self.dnn(t)
    u_t = self.dt(u, t)
    ode = u_t - self.renorm*(self.mu + (self.A @ torch.exp(u).T).T)
    return ode
</code></pre>
<p>the ODE system to be solved is ∂<em>t</em> <em>ui</em> = <em>µ</em> + <em>A</em>⋅exp(<em>ui</em>), where <em>µ</em> ∈ ℝ^N and <em>A</em> ∈ ℝ^(N,N) are stored in the matrix <code>theta</code>, are fixed parameters involved in the ODE.</p>
<p>How could I implement the function <code>dt</code> so that there is no more loop in it ?<br>
I also tried the decorator <code>jit</code> but without success either…</p>
<p>(<em>nb :</em> this is a duplicate of my stack overflow post that did not get an answer yet…)</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/ergias">@ergias</a>,</p>
<p>In the other discuss link you shared, it referenced my answer of using the <code>functorch</code> library, however, with the release of PyTorch2.0+, the <code>functorch</code> library is now within <code>torch</code> itself in the <code>torch.func</code> namespace.</p>
<p>For the most part things have stayed the same, except instead of functionalizing your <code>net</code> directly you need to call <code>torch.func.functional_call(net, params, *inputs)</code> and pass the parameters of your network as a dictionary.</p>
<p>For example,</p>
<pre><code class="lang-auto">net = Net(*args, **kwargs)
from torch.func import functional_call, jacrev, vmap

def fcall(params, x):
  return functional_call(net, params, x)

per_sample_jacobian_wrt_params = vmap(jacrev(fcall, argnums=(0)), in_dims=(None,0))(params, x)
</code></pre>
<p>If you want to remove your loop you’ll need to do a simpler trick what I showed above. If you have more outputs to your function than inputs, you may want to use forward-mode AD via <code>torch.func.jacfwd</code> to get better performance.</p>
<p>Also, don’t use the <code>torch.jit()</code> decorator it really won’t do much, it’s best to use <code>torch.compile()</code> instead, but this will only give minimal improvements as well as <code>torch.compile</code> was only released with PyTorch2.0 recently.</p>
<p>Furthermore, be careful with re-wrapping your tensor, e.g. <code>self.A</code> and <code>self.mu</code> as that’ll break any gradient flow for those parameters. You could also directly initialize on the GPU via <code>torch.as_tensor(data, dtype=torch.float, device=self.device)</code>, which may be slightly quicker too.</p>
<aside class="quote no-group" data-username="ergias" data-post="1" data-topic="185929">
<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/ergias/48/62994_2.png" class="avatar"> ergias:</div>
<blockquote>
<pre><code class="lang-auto">    self.A = torch.tensor(theta[:,1:], requires_grad=False).float().to(self.device)
    self.mu = torch.tensor(theta[:,0], requires_grad=False).float().to(self.device)
</code></pre>
</blockquote>
</aside>

ergias: <p>Hi <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a> , thanks for your answer !</p>
<p>However, I was not able to use the tips you gave in the code.<br>
If I have well understood :</p>
<ul>
<li>I need to convert the method <code>self.dt</code> into a <code>functional_call</code>, that takes the same arguments</li>
<li>then the call of <code>vmap</code> to this <code>functional_call</code> will return the gradient computed on each column, but more efficiently.</li>
</ul>
<p>Here is what I have done :</p>
<pre><code class="lang-py">def dt(self, u, t):
    nb_species = u.shape[1]
    u_t = [] # initializing the time derivatives
    
    # for each species (i), we compute the time derivative ui_t
    for i in range(nb_species):
        ui_t = torch.autograd.grad(u[:, i], t,
                                    grad_outputs=torch.ones_like(u[:, i]),
                                    retain_graph=True,
                                    create_graph=True,
                                    allow_unused=True)[0]
        u_t.append(ui_t) # time derivatives are stored in u_t
    u_t = torch.cat(tuple(u_t), dim=1) # we concateneate all the derivatives
    return u_t

def dt_call(self, u, t):
    return torch.func.functional_call(self.dt, u, t)

def functional_f(self, t, theta):
    self.A = torch.tensor(theta[:,1:], requires_grad=False).float().to(self.device)
    self.mu = torch.tensor(theta[:,0], requires_grad=False).float().to(self.device)

    u = self.dnn(t)
    u_t = torch.func.vmap(torch.func.jacrev(self.dt_call, argnums=(0)), in_dims=(None, 0))(u, t)
    # u_t = self.dt(u, t)
    ode = u_t - self.renorm*(self.mu + (self.A @ torch.exp(u).T).T)
    return ode
</code></pre>
<p>I did not really catch what was the argument <code>in_dims</code> argument of the function <code>vmap</code>, and I think it is the key point because I get this error : <code>ValueError: Expected parameter_and_buffer_dicts to be a dict, or a list/tuple of dicts, but got &lt;class 'torch.Tensor'&gt;</code> (I also tried to change the values of the argument, but it raises a similar error <code>ValueError: vmap: Expected all tensors to have the same size in the mapped dimension, got sizes [3, 21] for the mapped dimension</code>).</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/ergias">@ergias</a>,</p>
<p>In my example, the functional_call is applied to <code>nn.Module</code> objects, in your case you have just a function and you can simply pass it to <code>torch.func.vmap</code>. I’d recommend giving the docs for <code>torch.func.grad</code> and <code>torch.func.vmap</code> a read so you can apply them to your particular use-case, as you don’t have a minimal reproducible example it’s hard to see how everything fits together.</p>
<p>If you can share a <em>minimal</em> reproducible example, I can have a look.</p>