ProvenanceTensor bug when used with torch.set_default_device

[ordabayevy]

When trying to debug the code from the forum post I came across this bug:

import torch
from pyro.ops.provenance import ProvenanceTensor

device = torch.device("cuda")
torch.set_default_device(device)
x = torch.tensor([1., 2., 3.])
y = ProvenanceTensor(x, frozenset(["x"]))
print(torch.as_tensor(y))

returns tensor([], device='cuda:0')

Two observations:

1. This doesn't happen if torch.set_default_device is not used: print(torch.as_tensor(y.cuda())) works fine
2. When torch.set_default_device is used then the following code gets invoked
   https://github.com/pytorch/pytorch/blob/main/torch/utils/_device.py#L72-L76
   where func(*args, **kwargs) returns tensor([], device='cuda:0'). This doesn't happen when .cuda() is used like in 1).

[ilia-kats]

I don't think it's related to set_default_device, but rather more generally to device handling in as_tensor:

import torch
from pyro.ops.provenance import ProvenanceTensor

device = torch.device("cuda")
x = torch.tensor([1., 2., 3.])
y = ProvenanceTensor(x, frozenset(["x"]))
print(torch.as_tensor(y, device=device))

also returns an empty tensor

[ilia-kats]

I think device handing in ProvenanceTensor is generally buggy:

import torch
from pyro.ops.provenance import ProvenanceTensor

device = torch.device("cuda")
x = torch.tensor([1., 2., 3.])
y = ProvenanceTensor(x, frozenset(["x"]))
print(torch.as_tensor(y.cuda(), device="cpu").device)

returns device(type='cuda', index=0)

but

print(torch.as_tensor(torch.arange(5).cuda(), device="cpu").device)

returns device(type='cpu') as expected

[ilia-kats]

Apparently ProvenanceTensor.__torch_function__ is not being called in as_tensor. It is being called for y.cuda(), though