make ProvenanceTensor behave more like a Tensor (closes #3218)

[ilia-kats]

The data is now stored in the Tensor itself instead of an attribute. This fixes torch.to_tensor returning empty tensors when called with a ProvenanceTensor and and a device as arguments.

This is super hacky, but I couldn't come up with a cleaner way. Note that this is the only way to use pyro.infer.inspect.get_dependencies when training on GPUs (I'm using it in a custom Messenger guide), since the log_prob function of some distributions (for example Gamma) calls torch.to_tensor.

[fritzo]

Thanks for this subtle fix! This looks good to me (after one minor comment), but I'm unsure how this will interact with other subclasses of tensor.

@ordabayevy could you also take a look as you've thought about this before?

[fritzo]

Could this line be simplified to

ret._t = data

or would that break something?

[ilia-kats]

Right, thanks. Took me about four tries to get all the tests to pass, this was still a remnant of an earlier attempt.

[fritzo]

Would you be able to add a regression test and decorate it with @requires_cuda? It won't run on CI, but it might help future maintainers of ProvenanceTensor preserve your intent in the face of future changes.

[fritzo]

LGTM, thanks for adding a test.

I'll leave this up a couple days before merging in case @ordabayevy has any comments.

[ordabayevy]

Thanks for holding it up. I'll have a look at this later tonight.

[ordabayevy]

@ilia-kats thanks for fixing this!

What about trying to use torch.Tensor._make_subclass as is used by torch.nn.Parameter:

class ProvenanceTensor(torch.Tensor):
         assert not isinstance(data, ProvenanceTensor)
         if not provenance:
             return data
-        return super().__new__(cls)
+        return torch.Tensor._make_subclass(cls, data)

And I believe we can remove instance check from the __init__ method since data being a ProvenanceTensor is prohibited in the __new__ method anyways? WDYT @fritzo @ilia-kats ?

     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

[ilia-kats]

@ordabayevy Thanks for the comment. I actually Tensor.as_subclass in the public API, which also seems to work.

[fritzo]

@ordabayevy ready to merge? I'll release today or tomorrow and will include this PR in the release

[ordabayevy]

Yeah, lgtm.