`tensor` not a `FakeTensor` under `FakeTensorMode` and `device('meta')`

[carmocca]

🐛 Describe the bug

import torch
from torch._subclasses.fake_tensor import FakeTensorMode

with FakeTensorMode(), torch.device("meta"):
    a = torch.tensor(3.0)  # not a FakeTensor
    b = torch.full(size=tuple(), fill_value=3.0)  # workaround

print(a)  # tensor(..., device='meta', size=())
print(b)  # FakeTensor(..., device='meta', size=())

You would expect both to become FakeTensors

This issue results in a downstream assertion error

pytorch/torch/_subclasses/fake_tensor.py

Lines 2230 to 2233 in 01b055a

	raise AssertionError(
	f"Please convert all Tensors to FakeTensors first or instantiate FakeTensorMode "
	f"with 'allow_non_fake_inputs'. Found in {render_call(func, args, kwargs)}"
	)

in scripts that try to use this tech

Versions

torch==2.6.0.dev20241018+cu118

cc @ezyang @chauhang @penguinwu @eellison @zou3519 @bdhirsh @yf225

[ezyang]

Weird... what happens if you interchange the order of the context managers?

[carmocca]

Same output

[bdhirsh]

This doesn't repro for me if I take away device("meta"):

with FakeTensorMode():
    a = torch.tensor(3.0)  # now becomes a FakeTensor

[bdhirsh]

Looks like it's because when you call torch.tensor(3.0, device="meta") we have some logic to short-circuit in the meta case, which prevents us from going on to call into any modes: https://github.com/pytorch/pytorch/blob/main/torch/csrc/utils/tensor_new.cpp#L436

We should probably make sure that in the meta case, we still end up calling at::lift_fresh()

[yiming0416]

Doing old issue scrubbing, this still repros on main.
Is there any action items we can take? cc @bdhirsh

[abhitorch81]

I would like to explore the issue.