Repository: pytorch
Issue: torch.vmap + autograd.Function + current_level bug
Labels: module: autograd,triaged,module: vmap,module: functorch

### 🐛 Describe the bug

It we call ```torch._C._functorch.current_level()``` inside of an autograd function's ```setup_context``` method, and then apply ```torch.vmap``` on top of it, it errors out.
Repro:
* Checkout and apply https://github.com/pytorch/pytorch/pull/143811
* Replace ```key = id(Generated)``` with ```key = current_level()``` in ```setup_context```.
* Run the following example:
```
import torch

class LinearFunction(torch.autograd.Function):
    generate_vmap_rule = True

    # Note that forward, setup_context, and backward are @staticmethods
    @staticmethod
    def forward(input, weight, bias):
        output = input.mm(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output

    @staticmethod
    # inputs is a Tuple of all of the inputs passed to forward.
    # output is the output of the forward().
    def setup_context(ctx, inputs, output):
        input, weight, bias = inputs
        ctx.save_for_backward(input, weight, bias)

    # This function has only a single output, so it gets only one gradient
    @staticmethod
    def backward(ctx, grad_output):
        input, weight, bias = ctx.saved_tensors
        grad_input = grad_weight = grad_bias = None
        if ctx.needs_input_grad[0]:
            grad_input = grad_output.mm(weight)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().mm(input)
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias

def fn(input, weight, bias=None):
    return torch.vmap(LinearFunction.apply)(input, weight, bias)


torch.manual_seed(124)
batch_input = torch.randn(4, 2, 2, dtype=torch.double, requires_grad=True)
batch_weight = torch.randn(4, 3, 2, dtype=torch.double, requires_grad=True)
batch_bias = torch.randn(4, 3, dtype=torch.double, requires_grad=True)
output = fn(batch_input, batch_weight, batch_bias)
print(output)
```

Then it errors out:
```
Traceback (most recent call last):
  File "/data/users/ybliang/debug/debug7.py", line 44, in <module>
    output = fn(batch_input, batch_weight, batch_bias)
  File "/data/users/ybliang/debug/debug7.py", line 37, in fn
    return torch.vmap(LinearFunction.apply)(input, weight, bias)
  File "/home/ybliang/local/pytorch/torch/_functorch/apis.py", line 203, in wrapped
    return vmap_impl(
  File "/home/ybliang/local/pytorch/torch/_functorch/vmap.py", line 331, in vmap_impl
    return _flat_vmap(
  File "/home/ybliang/local/pytorch/torch/_functorch/vmap.py", line 481, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "/home/ybliang/local/pytorch/torch/autograd/function.py", line 585, in apply
    return custom_function_call(cls, *args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_functorch/autograd_function.py", line 49, in __call__
    return super().__call__(autograd_function, *args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_ops.py", line 439, in __call__
    return wrapper()
  File "/home/ybliang/local/pytorch/torch/_dynamo/eval_frame.py", line 755, in _fn
    return fn(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_ops.py", line 435, in wrapper
    return self.dispatch(
  File "/home/ybliang/local/pytorch/torch/_ops.py", line 305, in dispatch
    return dispatch_functorch(self, args, kwargs)
  File "/home/ybliang/local/pytorch/torch/_functorch/pyfunctorch.py", line 294, in dispatch_functorch
    return interpreter.process(op, args, kwargs)
  File "/home/ybliang/local/pytorch/torch/_functorch/pyfunctorch.py", line 130, in process
    return kernel(self, *args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/_functorch/autograd_function.py", line 300, in custom_function_call_vmap
    return custom_function_call_vmap_generate_rule(
  File "/home/ybliang/local/pytorch/torch/_functorch/autograd_function.py", line 384, in custom_function_call_vmap_generate_rule
    outputs = custom_function_call(vmapped_function, *unwrapped_operands)
  File "/home/ybliang/local/pytorch/torch/_functorch/autograd_function.py", line 50, in __call__
    return autograd_function.apply(*args, **kwargs)
  File "/home/ybliang/local/pytorch/torch/autograd/function.py", line 575, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/home/ybliang/local/pytorch/torch/_functorch/autograd_function.py", line 410, in setup_context
    key = current_level()
RuntimeError: maybe_layer.has_value() INTERNAL ASSERT FAILED at "/data/users/ybliang/pytorch/torch/csrc/functorch/init.cpp":370, please report a bug to PyTorch. 
```


### Versions

main

cc @ezyang @albanD @gqchen @pearu @nikitaved @soulitzer @Varal7 @xmfan @zou3519 @Chillee @samdow @kshitij12345