Repository: pyro-ppl/pyro

Issue:
`ProvenanceTensor` bug when used with `torch.set_default_device`
When trying to debug the code from the [forum post](https://forum.pyro.ai/t/running-pyro-plate-on-gpu/5097/8) I came across this bug:

```py
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).


Discussion:
I don't think it's related to `set_default_device`, but rather more generally to device handling in `as_tensor`:
```python
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
I think device handing in `ProvenanceTensor` is generally buggy:
```python
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
```python
print(torch.as_tensor(torch.arange(5).cuda(), device="cpu").device)
```
returns `device(type='cpu')` as expected
Apparently `ProvenanceTensor.__torch_function__` is not being called in `as_tensor`. It is being called for `y.cuda()`, though

Resolving patch:
diff --git a/pyro/ops/provenance.py b/pyro/ops/provenance.py
--- a/pyro/ops/provenance.py
+++ b/pyro/ops/provenance.py
@@ -46,14 +46,15 @@ def __new__(cls, data: torch.Tensor, provenance=frozenset(), **kwargs):
         assert not isinstance(data, ProvenanceTensor)
         if not provenance:
             return data
-        return super().__new__(cls)
+        ret = data.as_subclass(cls)
+        ret._t = data  # this makes sure that detach_provenance always
+        # returns the same object. This is important when
+        # using the tensor as key in a dict, e.g. the global
+        # param store
+        return ret
 
     def __init__(self, data, provenance=frozenset()):
         assert isinstance(provenance, frozenset)
-        if isinstance(data, ProvenanceTensor):
-            provenance |= data._provenance
-            data = data._t
-        self._t = data
         self._provenance = provenance
 
     def __repr__(self):


Test patch:
diff --git a/tests/ops/test_provenance.py b/tests/ops/test_provenance.py
new file mode 100644
--- /dev/null
+++ b/tests/ops/test_provenance.py
@@ -0,0 +1,45 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+
+from pyro.ops.provenance import ProvenanceTensor
+from tests.common import assert_equal, requires_cuda
+
+
+@requires_cuda
+@pytest.mark.parametrize(
+    "dtype1",
+    [
+        torch.float16,
+        torch.float32,
+        torch.float64,
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.uint8,
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype2",
+    [
+        torch.float16,
+        torch.float32,
+        torch.float64,
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.uint8,
+    ],
+)
+def test_provenance_tensor(dtype1, dtype2):
+    device = torch.device("cuda")
+    x = torch.tensor([1, 2, 3], dtype=dtype1)
+    y = ProvenanceTensor(x, frozenset(["x"]))
+    z = torch.as_tensor(y, device=device, dtype=dtype2)
+
+    assert x.shape == y.shape == z.shape
+    assert_equal(x, z.cpu())