🐛 Describe the bug

🐛 Describe the bug

torch.compile crashes when compiling a vmapped Hessian computation (torch.vmap(torch.func.hessian(f))). The crash occurs at the Dynamo level (not Inductor-specific — aot_eager backend also crashes).

The root cause appears to be that jacfwd as the outer transform, combined with vmap and compile, fails during fake tensor propagation. Specifically, hessian(f) = jacfwd(jacrev(f)), and the jacfwd outer layer triggers the crash.

Key observations:

• vmap(jacfwd(jacrev(f))) (= vmap(hessian(f))) — CRASHES
• vmap(jacfwd(jacfwd(f))) — CRASHES
• vmap(jacrev(jacrev(f))) — WORKS (workaround)
• vmap(jacrev(jacfwd(f))) — WORKS
• jacfwd(jacrev(f)) without vmap — WORKS
• vmap(jacfwd(f)) without nesting — WORKS

So the crash requires: jacfwd as outer + nested differentiation + vmap + compile.

Error message

TorchRuntimeError: RuntimeError when making fake tensor call
Explanation: Dynamo failed to run FX node with fake tensors:
  call_function <function _autograd_grad at 0x...>(
    *([GradTrackingTensor(lvl=-2, value=
        GradTrackingTensor(lvl=3, value=
          BatchedTensor(lvl=1, bdim=0, value=
            FakeTensor(..., device='cuda:0', size=(4,), dtype=torch.float64))))],
      ...))

To reproduce

import torch

torch.manual_seed(42)
x = torch.randn(4, 8, device='cuda', dtype=torch.float64)

def f(x):
    return (x ** 4).sum()

# CRASHES
compiled = torch.compile(torch.vmap(torch.func.hessian(f)), fullgraph=True)
result = compiled(x)

# WORKAROUND (works correctly)
compiled_workaround = torch.compile(
    torch.vmap(torch.func.jacrev(torch.func.jacrev(f))),
    fullgraph=True
)
result = compiled_workaround(x)

The crash also occurs with f(x) = sin(x).sum(), f(x) = exp(x).sum(), f(x) = (x**2).sum(), but NOT with f(x) = tanh(x).sum().

Versions

Versions

PyTorch: 2.13.0.dev20260501+cu126
Python: 3.11
CUDA: 12.6
GPU: Tesla T4

cc @chauhang @penguinwu @Chillee @samdow @kshitij12345