Nested forward mode AD not supported (jacobian)— what operations can trigger this?

k3in: <p>Hi all,</p>
<p>I’m trying to compute a Jacobian using <code>torch.autograd.functional.jacobian</code> with <code>strategy='forward-mode'</code>, but I keep getting this error:</p>
<blockquote>
<p>File “/my_code.py”, line 694, in J_func<br>
J_ce = torch.autograd.functional.jacobian(ce_func, x, vectorize=True, strategy=‘forward-mode’)<br>
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
File “/venv/lib/python3.11/site-packages/torch/autograd/functional.py”, line 671, in jacobian<br>
return _jacfwd(func, inputs, strict, vectorize)<br>
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
File “/venv/lib/python3.11/site-packages/torch/autograd/functional.py”, line 548, in _jacfwd<br>
outputs_before_split = _vmap(jvp)(tangents)<br>
^^^^^^^^^^^^^^^^^^^^<br>
File “/venv/lib/python3.11/site-packages/torch/_vmap_internals.py”, line 231, in wrapped<br>
batched_outputs = func(*batched_inputs)<br>
^^^^^^^^^^^^^^^^^^^^^<br>
File “/venv/lib/python3.11/site-packages/torch/autograd/functional.py”, line 527, in jvp<br>
with fwAD.dual_level():<br>
File “/venv/lib/python3.11/site-packages/torch/autograd/forward_ad.py”, line 210, in <strong>enter</strong><br>
return enter_dual_level()<br>
^^^^^^^^^^^^^^^^^^<br>
File “/venv/lib/python3.11/site-packages/torch/autograd/forward_ad.py”, line 34, in enter_dual_level<br>
new_level = torch._C._enter_dual_level()<br>
^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
RuntimeError: Nested forward mode AD is not supported at the moment</p>
</blockquote>
<p>Since this Jacobian is the ONLY derivative I am computing, I don’t understand why nested AD would be needed. Before calling  torch.autograd.functional.jacobian I detach the variable and also reset the gradients so I don’t think any second pass over the graph sgould be triggered.</p>
<p>What are some possible operations that would cause such an error?</p>
<p>(Apologies for not providing code - the function I take the jacobian of is a convolution of many functions with a lot of project specific code. )</p>

KFrank: <p>Hi k3in!</p>
<aside class="quote no-group quote-modified" data-username="k3in" data-post="1" data-topic="221902" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/22d042/48.png" class="avatar"> k3in:</div>
<blockquote>
<p>I’m trying to compute a Jacobian using <code>torch.autograd.functional.jacobian</code> with <code>strategy='forward-mode'</code>, but I keep getting this error:<br>
…</p>
<blockquote>
<p>RuntimeError: Nested forward mode AD is not supported at the moment</p>
</blockquote>
</blockquote>
</aside>
<p>The simple explanation is that you are indeed applying forward-mode AD to the<br>
result of forward-mode AD.</p>
<p>Consider the following script that produces a stack-trace similar to yours:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

def sq (t):
    return  t*t

jacsqrev = torch.func.jacrev (sq)   # function object that evaluates the jacobian of sq at some argument
jacsqfwd = torch.func.jacfwd (sq)   # same, but uses forward mode internally

# compute first-order derivative
jacA = jacsqrev (torch.ones (1))
jacB = jacsqfwd (torch.ones (1))
jacC = torch.autograd.functional.jacobian (sq, torch.ones (1), vectorize = True, strategy = 'forward-mode')

# all three give the same result
print ('jacA:', jacA)
print ('jacB:', jacB)
print ('jacC:', jacC)

# compute second-order derivative

# forward-mode applied to reverse-mode works
jacjacA = torch.autograd.functional.jacobian (jacsqrev, torch.ones (1), vectorize = True, strategy = 'forward-mode')
# reverse-mode applied to forward-mode works
jacjacB = torch.autograd.functional.jacobian (jacsqfwd, torch.ones (1), vectorize = True, strategy = 'reverse-mode')

# both give the same result
print ('jacjacA:', jacjacA)
print ('jacjacB:', jacjacB)

# forward-mode applied to forward-mode will fail
jacjacC = torch.autograd.functional.jacobian (jacsqfwd, torch.ones (1), vectorize = True, strategy = 'forward-mode')
</code></pre>
<p>Here is its output:</p>
<pre><code class="lang-plaintext">2.7.1+cu128
jacA: tensor([[2.]])
jacB: tensor([[2.]])
jacC: tensor([[2.]])
jacjacA: tensor([[[2.]]])
jacjacB: tensor([[[2.]]])
Traceback (most recent call last):
  File "&lt;python-input-4&gt;", line 1, in &lt;module&gt;
    exec (open ('./nested_fwd.py').read())
    ~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "&lt;string&gt;", line 36, in &lt;module&gt;
  File "&lt;path_to_pytorch_install&gt;/torch/autograd/functional.py", line 670, in jacobian
    return _jacfwd(func, inputs, strict, vectorize)
  File "&lt;path_to_pytorch_install&gt;/torch/autograd/functional.py", line 547, in _jacfwd
    outputs_before_split = _vmap(jvp)(tangents)
  File "&lt;path_to_pytorch_install&gt;/torch/_vmap_internals.py", line 231, in wrapped
    batched_outputs = func(*batched_inputs)
  File "&lt;path_to_pytorch_install&gt;/torch/autograd/functional.py", line 532, in jvp
    func(*dual_inputs), "outputs"
    ~~~~^^^^^^^^^^^^^^
  File "&lt;path_to_pytorch_install&gt;/torch/_functorch/eager_transforms.py", line 1273, in wrapper_fn
    results = vmap(push_jvp, randomness=randomness)(basis)
  File "&lt;path_to_pytorch_install&gt;/torch/_functorch/apis.py", line 202, in wrapped
    return vmap_impl(
        func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs
    )
  File "&lt;path_to_pytorch_install&gt;/torch/_functorch/vmap.py", line 334, in vmap_impl
    return _flat_vmap(
        func,
    ...&lt;6 lines&gt;...
        **kwargs,
    )
  File "&lt;path_to_pytorch_install&gt;/torch/_functorch/vmap.py", line 484, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "&lt;path_to_pytorch_install&gt;/torch/_functorch/eager_transforms.py", line 1262, in push_jvp
    output = _jvp_with_argnums(
        func, args, basis, argnums=argnums, has_aux=has_aux
    )
  File "&lt;path_to_pytorch_install&gt;/torch/_functorch/eager_transforms.py", line 1093, in _jvp_with_argnums
    with ctx():
         ~~~^^
  File "&lt;path_to_pytorch_install&gt;/torch/autograd/forward_ad.py", line 210, in __enter__
    return enter_dual_level()
  File "&lt;path_to_pytorch_install&gt;/torch/autograd/forward_ad.py", line 34, in enter_dual_level
    new_level = torch._C._enter_dual_level()
RuntimeError: Nested forward mode AD is not supported at the moment
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>Since this Jacobian is the ONLY derivative I am computing, I don’t understand why nested AD would be needed.<br>
…<br>
What are some possible operations that would cause such an error?</p>
</blockquote>
</aside>
<p>If it’s really true that you’re not running forward-mode AD twice, then the only thing<br>
I can image is that maybe your complicated, secret code is setting up some dual<br>
variables, but not using them for AD (but I haven’t tried this to check the results).</p>
<p>Can you perform divide-and-conquer debugging on your complicated, secret code?<br>
I.e., block out half of the code to which you apply forward-mode <code>jacobian</code> and then<br>
the other half.  Can you make the error go away?  Do this again on the section of your<br>
code the triggers the error.  At some point you should whittle things down to a short<br>
fragment of code that triggers the error.</p>
<p>Good luck!</p>
<p>K. Frank</p>