import torch
import functorch as ft

def foo(t):
    tt = t.detach()
    n = tt.numpy()
    return t

x = torch.rand(4, 3)
out = ft.grad(foo)(x)
# or
# out = ft.vmap(foo)(x)

# Monkeypatch .numpy() to fetch underlying tensor and call .numpy()
_old_numpy = torch.Tensor.numpy


@functools.wraps(_old_numpy)
def _numpy(tensor):
    level = _C.maybe_get_level(tensor)
    if level == -1:
        return _old_numpy(tensor)

    if _C.is_functionaltensor(tensor):
        # Since we're unwrapping the FunctionalTensorWrapper, we need to make sure
        # that it's up to date first
        torch._sync(tensor)

    value = _C.get_unwrapped(tensor)
    dl_enabled = _C.tls_set_is_included()
    try:
        # Disable temporarily kDynamicLayerFrontModeKey/kDynamicLayerBackModeKey as included dispatch keys
        if (dl_enabled):
            _C._set_dynamic_layer_keys_included(False)
        return value.numpy()
    finally:
        # Reenable kDynamicLayerFrontModeKey/kDynamicLayerBackModeKey as included dispatch keys
        if (dl_enabled):
            _C._set_dynamic_layer_keys_included(True)


setattr(torch.Tensor, 'numpy', _numpy)

import torch
import functorch as ft

def foo(t):
    n = t.numpy()
    assert n.shape == (4, 3)
    assert n.shape != (3, )
    return t

x = torch.rand(4, 3)
out = ft.vmap(foo)(x)