ProvenanceTensor

Makefile

@@ -53,13 +53,16 @@ ifeq (${FUNSOR_BACKEND}, torch)
 	python examples/adam.py --num-steps=21
 	@echo PASS
 else ifeq (${FUNSOR_BACKEND}, jax)
-	pytest -v -n auto --ignore=test/examples --ignore=test/pyro --ignore=test/pyroapi --ignore=test/test_distribution.py --ignore=test/test_distribution_generic.py
+	pytest -v -n auto --ignore=test/examples --ignore=test/pyro --ignore=test/pyroapi \
+		--ignore=test/test_distribution.py --ignore=test/test_distribution_generic.py \
+		--ignore=test/torch
 	pytest -v -n auto test/test_distribution.py
 	pytest -v -n auto test/test_distribution_generic.py
 	@echo PASS
 else
 	# default backend
-	pytest -v -n auto --ignore=test/examples --ignore=test/pyro --ignore=test/pyroapi
+	pytest -v -n auto --ignore=test/examples --ignore=test/pyro \
+		--ignore=test/pyroapi --ignore=test/torch
 	@echo PASS
 endif
 

funsor/torch/provenance.py

@@ -0,0 +1,65 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+
+class ProvenanceTensor(torch.Tensor):
+    """
+    Provenance tracking implementation in Pytorch.
+
+    Provenance of the output tensor is the union of provenances of input tensors.
+    """
+
+    def __new__(cls, data, provenance=frozenset(), **kwargs):
+        if not provenance:
+            return data
+        instance = torch.Tensor.__new__(cls)
+        instance.__init__(data, provenance)
+        return instance
+
+    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):
+        return "Provenance:\n{}\nTensor:\n{}".format(self._provenance, self._t)
+
+    def __torch_function__(self, func, types, args=(), kwargs=None):
+        if kwargs is None:
+            kwargs = {}
+        # collect provenance information from args
+        provenance = frozenset()
+        # extract ProvenanceTensor._t data from args
+        _args = []
+        for arg in args:
+            if isinstance(arg, ProvenanceTensor):
+                provenance |= arg._provenance
+                _args.append(arg._t)
+            elif isinstance(arg, tuple):
+                _arg = []
+                for a in arg:
+                    if isinstance(a, ProvenanceTensor):
+                        provenance |= a._provenance
+                        _arg.append(a._t)
+                    else:
+                        _arg.append(a)
+                _args.append(tuple(_arg))
+            else:
+                _args.append(arg)
+        ret = func(*_args, **kwargs)
+        if isinstance(ret, torch.Tensor):
+            return ProvenanceTensor(ret, provenance=provenance)
+        if isinstance(ret, tuple):
+            _ret = []
+            for r in ret:
+                if isinstance(r, torch.Tensor):
+                    _ret.append(ProvenanceTensor(r, provenance=provenance))
+                else:
+                    _ret.append(r)
+            return tuple(_ret)
+        return ret

test/torch/test_provenance.py

@@ -0,0 +1,103 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+from pyro.ops.indexing import Vindex
+
+from funsor.torch.provenance import ProvenanceTensor
+
+
+@pytest.mark.parametrize("op", ["log", "exp", "long"])
+@pytest.mark.parametrize(
+    "data,provenance",
+    [
+        (torch.tensor([1]), "ab"),
+        (torch.tensor([1]), "a"),
+    ],
+)
+def test_unary(op, data, provenance):
+    data = ProvenanceTensor(data, frozenset(provenance))
+
+    expected = frozenset(provenance)
+    actual = getattr(data, op)()._provenance
+    assert actual == expected
+
+
+@pytest.mark.parametrize("data1,provenance1", [(torch.tensor([1]), "a")])
+@pytest.mark.parametrize(
+    "data2,provenance2",
+    [
+        (torch.tensor([2]), "b"),
+        (torch.tensor([2]), ""),
+        (2, ""),
+    ],
+)
+def test_binary_add(data1, provenance1, data2, provenance2):
+    data1 = ProvenanceTensor(data1, frozenset(provenance1))
+    if provenance2:
+        data2 = ProvenanceTensor(data2, frozenset(provenance2))
+
+    expected = frozenset(provenance1 + provenance2)
+    actual = torch.add(data1, data2)._provenance
+    assert actual == expected
+
+
+@pytest.mark.parametrize(
+    "data1,provenance1",
+    [
+        (torch.tensor([0, 1]), "a"),
+        (torch.tensor([0, 1]), ""),
+    ],
+)
+@pytest.mark.parametrize(
+    "data2,provenance2",
+    [
+        (torch.tensor([0]), "b"),
+        (torch.tensor([1]), ""),
+    ],
+)
+def test_indexing(data1, provenance1, data2, provenance2):
+    if provenance1:
+        data1 = ProvenanceTensor(data1, frozenset(provenance1))
+    if provenance2:
+        data2 = ProvenanceTensor(data2, frozenset(provenance2))
+
+    expected = frozenset(provenance1 + provenance2)
+    actual = getattr(data1[data2], "_provenance", frozenset())
+    assert actual == expected
+
+
+@pytest.mark.parametrize(
+    "data1,provenance1",
+    [
+        (torch.tensor([[0, 1], [2, 3]]), "a"),
+        (torch.tensor([[0, 1], [2, 3]]), ""),
+    ],
+)
+@pytest.mark.parametrize(
+    "data2,provenance2",
+    [
+        (torch.tensor([0.0, 1.0]), "b"),
+        (torch.tensor([0.0, 1.0]), ""),
+    ],
+)
+@pytest.mark.parametrize(
+    "data3,provenance3",
+    [
+        (torch.tensor([0, 1]), "c"),
+        (torch.tensor([0, 1]), ""),
+    ],
+)
+def test_vindex(data1, provenance1, data2, provenance2, data3, provenance3):
+    if provenance1:
+        data1 = ProvenanceTensor(data1, frozenset(provenance1))
+    if provenance2:
+        data2 = ProvenanceTensor(data2, frozenset(provenance2))
+    if provenance3:
+        data3 = ProvenanceTensor(data3, frozenset(provenance3))
+
+    expected = frozenset(provenance1 + provenance2 + provenance3)
+    result = Vindex(data1)[data2.long().unsqueeze(-1), data3]
+    actual = getattr(result, "_provenance", frozenset())
+    assert actual == expected