Implementation:Pyro ppl Pyro ProvenanceTensor

Property	Value
Module	pyro.ops.provenance
Source	pyro/ops/provenance.py
Lines	175
Classes	ProvenanceTensor
Functions	track_provenance, extract_provenance, get_provenance, detach_provenance
Dependencies	torch, torch.utils._pytree

Overview

This module implements provenance tracking for PyTorch tensors, enabling the framework to determine which input variables influenced which outputs through a computation graph. Provenance is represented as a frozenset of user-defined objects that propagates through PyTorch operations -- the provenance of any output tensor is the union of the provenances of its input tensors.

This is a form of nonstandard interpretation of programs, inspired by the work of Wingate et al. (2011) on efficient inference. In Pyro, provenance tracking is used to analyze dependencies between random variables, enabling more efficient inference algorithms.

The ProvenanceTensor class extends torch.Tensor using the __torch_function__ protocol, automatically propagating provenance through all standard PyTorch operations without requiring explicit tracking code.

Code Reference

Class: ProvenanceTensor

A torch.Tensor subclass that carries provenance metadata.

• __new__(cls, data, provenance): Creates a new ProvenanceTensor. If provenance is empty, returns the raw tensor.
• __torch_function__(cls, func, types, args, kwargs): Intercepts all PyTorch operations, strips provenance from inputs, executes the operation, then re-attaches the union of all input provenances to outputs.
• _t: Stores the underlying unwrapped tensor (important for identity-based dict keys like the param store).
• _provenance: The frozenset of provenance labels.

Function: track_provenance(x, provenance)

Adds provenance to tensor leaves of a data structure. Dispatches via singledispatch:

• torch.Tensor: Wraps in ProvenanceTensor
• ProvenanceTensor: Merges provenances
• Lists, tuples, dicts: Recursively applies via tree_map
• Sets, frozensets: Recursively applies to elements

Function: extract_provenance(x)

Separates a data structure into a detached value and its provenance. Returns (value, provenance) tuple.

Function: get_provenance(x)

Returns just the provenance frozenset of a data structure.

Function: detach_provenance(x)

Strips provenance from a data structure, analogous to torch.Tensor.detach() for gradients. Returns the unwrapped value.

I/O Contract

Function	Input	Output
ProvenanceTensor(data, provenance)	Tensor, frozenset	ProvenanceTensor (or raw Tensor if provenance is empty)
track_provenance(x, provenance)	Any (nested structure with Tensor leaves), frozenset	Same structure with ProvenanceTensor leaves
extract_provenance(x)	Any (possibly containing ProvenanceTensors)	Tuple (detached_value, provenance: frozenset)
get_provenance(x)	Any	frozenset
detach_provenance(x)	Any	Same structure with raw Tensor leaves

Usage Examples

import torch
from pyro.ops.provenance import (
    ProvenanceTensor,
    get_provenance,
    detach_provenance,
)

# Track provenance of tensors
a = ProvenanceTensor(torch.randn(3), frozenset({"a"}))
b = ProvenanceTensor(torch.randn(3), frozenset({"b"}))
c = torch.randn(3)  # no provenance

# Provenance propagates through operations
result = a + b + c
assert get_provenance(result) == frozenset({"a", "b"})

# Detach provenance to block propagation
result2 = a + detach_provenance(b) + c
assert get_provenance(result2) == frozenset({"a"})

# Works with nested data structures
data = {"x": a, "y": [b, c]}
prov = get_provenance(data)
assert prov == frozenset({"a", "b"})

Related Pages

• Pyro_ppl_Pyro_Util -- General Pyro utility functions