Incorrect behavior on vmap(torch.autograd.functional.jacobian)

[zou3519]

Repro (uses facebookresearch/theseus): https://gist.github.com/zou3519/b645379a7afa56a170c9c49a0b9e2138

There is silently incorrect behavior on vmap over torch.autograd.functional.jacobian, but it works out if we vmap over functorch.jacrev instead.

This should either raise an error or not be silently incorrect.

[zou3519]

Repro without theseus:

from functorch import vmap, vjp, jacrev
import torch

def f(x):
    theta = torch.linalg.norm(x, dim=0, keepdim=True)
    return theta

USE_AUTOGRAD_FUNCTIONAL = True

def compute_jac_f(xi):
    if USE_AUTOGRAD_FUNCTIONAL:
        jacs = torch.autograd.functional.jacobian(f, xi)
        return jacs[0]
    else:
        jac = jacrev(f)(xi)
    return jac

xi = torch.rand(5,3)
result = vmap(compute_jac_f)(xi)
expected = torch.stack([compute_jac_f(xi[n]) for n in range(xi.shape[0])])

print(result)
print(expected)
assert torch.allclose(result, expected)