Hi k3in!
The simple explanation is that you are indeed applying forward-mode AD to the
result of forward-mode AD.
Consider the following script that produces a stack-trace similar to yours:
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')
Here is its output:
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 "<python-input-4>", line 1, in <module>
exec (open ('./nested_fwd.py').read())
~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "<string>", line 36, in <module>
File "<path_to_pytorch_install>/torch/autograd/functional.py", line 670, in jacobian
return _jacfwd(func, inputs, strict, vectorize)
File "<path_to_pytorch_install>/torch/autograd/functional.py", line 547, in _jacfwd
outputs_before_split = _vmap(jvp)(tangents)
File "<path_to_pytorch_install>/torch/_vmap_internals.py", line 231, in wrapped
batched_outputs = func(*batched_inputs)
File "<path_to_pytorch_install>/torch/autograd/functional.py", line 532, in jvp
func(*dual_inputs), "outputs"
~~~~^^^^^^^^^^^^^^
File "<path_to_pytorch_install>/torch/_functorch/eager_transforms.py", line 1273, in wrapper_fn
results = vmap(push_jvp, randomness=randomness)(basis)
File "<path_to_pytorch_install>/torch/_functorch/apis.py", line 202, in wrapped
return vmap_impl(
func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs
)
File "<path_to_pytorch_install>/torch/_functorch/vmap.py", line 334, in vmap_impl
return _flat_vmap(
func,
...<6 lines>...
**kwargs,
)
File "<path_to_pytorch_install>/torch/_functorch/vmap.py", line 484, in _flat_vmap
batched_outputs = func(*batched_inputs, **kwargs)
File "<path_to_pytorch_install>/torch/_functorch/eager_transforms.py", line 1262, in push_jvp
output = _jvp_with_argnums(
func, args, basis, argnums=argnums, has_aux=has_aux
)
File "<path_to_pytorch_install>/torch/_functorch/eager_transforms.py", line 1093, in _jvp_with_argnums
with ctx():
~~~^^
File "<path_to_pytorch_install>/torch/autograd/forward_ad.py", line 210, in __enter__
return enter_dual_level()
File "<path_to_pytorch_install>/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
If it’s really true that you’re not running forward-mode AD twice, then the only thing
I can image is that maybe your complicated, secret code is setting up some dual
variables, but not using them for AD (but I haven’t tried this to check the results).
Can you perform divide-and-conquer debugging on your complicated, secret code?
I.e., block out half of the code to which you apply forward-mode jacobian and then
the other half. Can you make the error go away? Do this again on the section of your
code the triggers the error. At some point you should whittle things down to a short
fragment of code that triggers the error.
Good luck!
K. Frank