Missing argument create_graph in torch.func api

LuoXin-s: <p>Hi, I want replace <code>torch.autograd.functional.jvp</code>/<code>torch.autograd.functional.vjp</code> with <code>torch.func.jvp</code>/<code>torch.func.vjp</code>. Since <code>torch.func.jvp</code> is more efficient compared to autograd ops.</p>
<p>However, I notice these apis are not equivalent, since I calculate jvp to create a loss, thus second-order gradient is required, and we have to set <code>create_graph=True</code>. There is no <code>create_graph</code> argument in <code>torch.func</code> api, so is <code>torch.func</code> does not intent to use in my case or it support higher-order gradient by default?</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/luoxin-s">@LuoXin-s</a>,</p>
<p>You need to compose <code>jacrev</code> with <code>jacrev</code> to create a 2nd-order derivative within the <code>torch.func</code> framework. The gradients are computed in a functional manner (so a more direct mathematical approach) within the <code>torch.func</code> namespace.</p>
<p>You can use <code>torch.func.hessian</code> directly,</p>
<pre><code class="lang-auto">import torch
from torch.func import hessian, vmap, functional_call

net = Model(*args, **kwargs) #our model

x = torch.randn(batch_size, num_inputs) #inputs
params = dict(net.named_parameters()) #params (need to be an input for torch.func to work)

#functionalize version (params are now an input to the model)
def fcall(parms, x):
  return functional_call(net, params, x)

def compute_hessian(params, x):
  _hessian_wrt_x = hessian(fcall, argnums=1)(params, x) #forward-over-reverse hessian calc.
  return _hessian_wrt_x

#We define a laplacian func for a single function, then vectorize over the batch.
laplacian = vmap(compute_hessian, in_dims=(None,0))(params, x) 
</code></pre>
<p>Or, you can compose <code>jacrev</code> with itself, which you can view as an equivalent of <code>create_graph</code> arg within the <code>torch.autograd</code> API (albeit that the graph is constructed directly).</p>
<pre><code class="lang-auto">def compute_hessian(params, x):
  _hessian_wrt_x = jacrev(jacrev(fcall, argnums=1), argnums=1)(params, x)
  return _hessian_wrt_x
</code></pre>
<p>And then use <code>torch.func.vmap</code> to vectorize over a set of samples, like in the first example above.</p>
<p>Here are some other examples on the forums:</p><aside class="quote quote-modified" data-post="6" data-topic="99644">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/implementing-calculation-of-the-laplacian/99644/6">Implementing calculation of the Laplacian</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" style="--category-badge-color: #AB9364; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    Hi <a class="mention" href="/u/rfmiotto">@rfmiotto</a>, 
You can do this quite efficiently within the torch.func namespace, which requires version torch2.1+. Here’s an example below for the laplacian with respect to an input x below, 
import torch
from torch.func import hessian, vmap, functional_call

net = Model(*args, **kwargs) #our model

x = torch.randn(batch_size, num_inputs) #inputs
params = dict(net.named_parameters()) #params (need to be an input for torch.func to work)

#functionalize version (params are now an input to the mod…
  </blockquote>
</aside>

<aside class="quote quote-modified" data-post="12" data-topic="80771">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/computing-batch-jacobian-efficiently/80771/12">Computing batch Jacobian efficiently</a> 
  </div>
  <blockquote>
    Just an update for anyone who reads the thread in the future, as of PyTorch2,the functorch library is now included in pytorch. So you can replace functorch with torch.func, for the most part the syntax is the same except if you have an nn.Module you’ll need to create a ‘functional’ version of your model. 
For example, 
model = myModel(*args, **kwargs) #our network

from torch.func import vmap, jacrev, functional_call

params = dict(model.named_parameters())
inputs = torch.randn(batch_size, input…
  </blockquote>
</aside>


LuoXin-s: <p>Thank you for your suggestion. However, I’m still confused. Let me explain further.</p>
<p>Actually, I want to calculate the following loss:</p>
<pre data-code-wrap="python"><code class="lang-python">def isometry_loss(self, input, func):
    u = torch.randn_like(input, device=input.device)
    Ju = torch.autograd.functional.jvp(func, input, u, create_graph=True)[1]
    JTJu = torch.autograd.functional.vjp(func, input, Ju, create_graph=True)[1]
    
    TrR = torch.sum(Ju.float() ** 2, dim=tuple(range(1, Ju.dim()))).mean()
    TrR2 = torch.sum(JTJu.float() ** 2, dim=tuple(range(1, JTJu.dim()))).mean()
    
    isometry_loss = TrR2 / TrR ** 2
    return isometry_loss
</code></pre>
<p>Here, <code>func</code> is my model and <code>input</code> is the model input. This loss serves as part of the overall loss (<code>loss += isometry_loss</code>) for backpropagation. Therefore, I have to rely on <code>loss.backward()</code> for backpropagation and parameter updating.</p>
<p>I’m actually not want to use the functionalized API, I just want to use fast JVP implementation.</p>

LuoXin-s: <p>I don’t how to convert it to the isometry loss in my case</p>

LuoXin-s: <p>I made some simplifications, and I found that the following losses have different gradients with respect to the network parameters:</p>
<pre data-code-wrap="python"><code class="lang-python">def iso_loss1():
    Ju = torch.autograd.functional.jvp(func, input, u, create_graph=True)[1]
    TrR = torch.sum(Ju.float() ** 2, dim=tuple(range(1, Ju.dim()))).mean()

    isometry_loss = TrR
    return isometry_loss

def iso_loss2():
    Ju = torch.func.jvp(func, (input,), (u,))[1]
    TrR = torch.sum(Ju.float() ** 2, dim=tuple(range(1, Ju.dim()))).mean()

    isometry_loss = TrR
    return isometry_loss
</code></pre>

AlphaBetaGamma96: <p>Why do you keep <code>create_graph==True</code> in the <code>iso_loss1</code>?</p>
<p>Can you share what the output of the two losses are?</p>

LuoXin-s: <p>I aim to incorporate isometry loss into my model’s optimization objective to enhance its isometric properties. The computation of isometry loss requires Jacobian-vector products (JVP), and it’s crucial to set <code>create_graph=True</code> to enable gradient backpropagation through the JVP calculation.</p>