Refactor ProvenanceTensor to use pytree

pyro/ops/provenance.py

@@ -1,10 +1,11 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
-from functools import singledispatch
+from functools import partial, singledispatch
 from typing import Tuple
 
 import torch
+from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
 
 
 class ProvenanceTensor(torch.Tensor):
@@ -51,35 +52,17 @@ def __new__(cls, data: torch.Tensor, provenance=frozenset(), **kwargs):
         # returns the same object. This is important when
         # using the tensor as key in a dict, e.g. the global
         # param store
+        ret._provenance = provenance
         return ret
 
-    def __init__(self, data, provenance=frozenset()):
-        assert isinstance(provenance, frozenset)
-        self._provenance = provenance
-
     def __repr__(self):
         return "Provenance:\n{}\nTensor:\n{}".format(self._provenance, self._t)
 
     @classmethod
     def __torch_function__(cls, func, types, args=(), kwargs=None):
-        if kwargs is None:
-            kwargs = {}
-        # collect provenance information from args
-        provenance = frozenset()
-        # extract ProvenanceTensor._t data from args and kwargs
-        _args = []
-        for arg in args:
-            _arg, _provenance = extract_provenance(arg)
-            _args.append(_arg)
-            provenance |= _provenance
-        _kwargs = {}
-        for k, v in kwargs.items():
-            _v, _provenance = extract_provenance(v)
-            _kwargs[k] = _v
-            provenance |= provenance
+        _args, _kwargs = detach_provenance([args, kwargs or {}])
         ret = func(*_args, **_kwargs)
-        _ret = track_provenance(ret, provenance)
-        return _ret
+        return track_provenance(ret, get_provenance([args, kwargs]))
 
 
 @singledispatch
@@ -98,13 +81,18 @@ def track_provenance(x, provenance: frozenset):
 
 
 @track_provenance.register(frozenset)
-@track_provenance.register(list)
 @track_provenance.register(set)
-@track_provenance.register(tuple)
-def _track_provenance_list(x, provenance: frozenset):
+def _track_provenance_set(x, provenance: frozenset):
     return type(x)(track_provenance(part, provenance) for part in x)
 
 
+@track_provenance.register(list)
+@track_provenance.register(tuple)
+@track_provenance.register(dict)
+def _track_provenance_pytree(x, provenance: frozenset):
+    return tree_map(partial(track_provenance, provenance=provenance), x)
+
+
 @track_provenance.register
 def _track_provenance_provenancetensor(x: ProvenanceTensor, provenance: frozenset):
     x_value, old_provenance = extract_provenance(x)
@@ -131,10 +119,8 @@ def _extract_provenance_tensor(x):
 
 
 @extract_provenance.register(frozenset)
-@extract_provenance.register(list)
 @extract_provenance.register(set)
-@extract_provenance.register(tuple)
-def _extract_provenance_list(x):
+def _extract_provenance_set(x):
     provenance = frozenset()
     values = []
     for part in x:
@@ -145,6 +131,19 @@ def _extract_provenance_list(x):
     return value, provenance
 
 
+@extract_provenance.register(list)
+@extract_provenance.register(tuple)
+@extract_provenance.register(dict)
+def _extract_provenance_pytree(x):
+    flat_args, spec = tree_flatten(x)
+    xs = []
+    provenance = frozenset()
+    for x, p in map(extract_provenance, flat_args):
+        xs.append(x)
+        provenance |= p
+    return tree_unflatten(xs, spec), provenance
+
+
 def get_provenance(x) -> frozenset:
     """
     Reads the provenance of a recursive datastructure possibly containing

tests/ops/test_provenance.py

@@ -4,7 +4,7 @@
 import pytest
 import torch
 
-from pyro.ops.provenance import ProvenanceTensor
+from pyro.ops.provenance import ProvenanceTensor, get_provenance, track_provenance
 from tests.common import assert_equal, requires_cuda
 
 
@@ -43,3 +43,26 @@ def test_provenance_tensor(dtype1, dtype2):
 
     assert x.shape == y.shape == z.shape
     assert_equal(x, z.cpu())
+
+
+@pytest.mark.parametrize(
+    "x",
+    [
+        torch.tensor([1, 2, 3]),
+        track_provenance(torch.tensor([1, 2, 3]), frozenset("y")),
+        frozenset([torch.tensor([0, 1]), torch.tensor([2, 3])]),
+        set([torch.tensor([0, 1]), torch.tensor([2, 3])]),
+        [torch.tensor([0, 1]), torch.tensor([2, 3])],
+        (torch.tensor([0, 1]), torch.tensor([2, 3])),
+        {"a": torch.tensor([0, 1]), "b": torch.tensor([2, 3])},
+        {
+            "a": track_provenance(torch.tensor([0, 1]), frozenset("y")),
+            "b": [torch.tensor([2, 3]), torch.tensor([4, 5])],
+        },
+    ],
+)
+def test_track_provenance(x):
+    new_provenance = frozenset("x")
+    old_provenance = get_provenance(x)
+    provenance = old_provenance | new_provenance
+    assert provenance == get_provenance(track_provenance(x, new_provenance))