From 1f42162a45b76f8701db6adc436bf959a565402a Mon Sep 17 00:00:00 2001 From: Mohit Arvind Khakharia Date: Tue, 4 Aug 2026 02:15:37 -0400 Subject: [PATCH] Preserve torch.Size through provenance tracking torch.Size subclasses tuple, so it was being handled by the list/tuple/dict pytree branches of track_provenance and extract_provenance. tree_flatten decomposes it into its plain int leaves and tree_unflatten rebuilds it as an ordinary tuple, so every torch.Size that passed through the provenance machinery came back out with its type erased. That matters because torch APIs overload on the distinction. Tensor.new() reads a torch.Size as a shape but a tuple as data, so inside ProvenanceTensor.__torch_function__ the detach_provenance() call turned Multinomial.sample()'s counts = samples.new(self._extended_shape(...)) into a 1-element tensor instead of a correctly shaped one, and the following scatter_add_ raised "index 2 is out of bounds for dimension 0 with size 1". Register torch.Size on both singledispatch functions and mark it as a leaf in the two pytree helpers so nested occurrences (args/kwargs) survive as well. --- pyro/ops/provenance.py | 25 +++++++++++++-- tests/ops/test_provenance.py | 61 +++++++++++++++++++++++++++++++++++- 2 files changed, 83 insertions(+), 3 deletions(-) diff --git a/pyro/ops/provenance.py b/pyro/ops/provenance.py index ff5e64ad7f..bf67d0d15a 100644 --- a/pyro/ops/provenance.py +++ b/pyro/ops/provenance.py @@ -10,6 +10,13 @@ _Tensor = TypeVar("_Tensor", bound=torch.Tensor) +def _is_size(x: object) -> bool: + # ``torch.Size`` subclasses ``tuple``, so the pytree helpers below would + # otherwise flatten it and rebuild it as a plain tuple. Treating it as a + # leaf keeps the type intact. + return isinstance(x, torch.Size) + + class ProvenanceTensor(torch.Tensor): """ Provenance tracking implementation in Pytorch. @@ -85,6 +92,13 @@ def track_provenance(x, provenance: frozenset): track_provenance.register(torch.Tensor)(ProvenanceTensor) +@track_provenance.register(torch.Size) +def _track_provenance_size(x, provenance: frozenset): + # A torch.Size holds plain ints, so there is no provenance to add; return + # it unchanged rather than letting the tuple handler rebuild it. + return x + + @track_provenance.register(frozenset) @track_provenance.register(set) def _track_provenance_set(x, provenance: frozenset): @@ -95,7 +109,9 @@ def _track_provenance_set(x, provenance: frozenset): @track_provenance.register(tuple) @track_provenance.register(dict) def _track_provenance_pytree(x, provenance: frozenset): - return tree_map(partial(track_provenance, provenance=provenance), x) + return tree_map( + partial(track_provenance, provenance=provenance), x, is_leaf=_is_size + ) @track_provenance.register @@ -123,6 +139,11 @@ def _extract_provenance_tensor(x): return x._t, x._provenance +@extract_provenance.register(torch.Size) +def _extract_provenance_size(x): + return x, frozenset() + + @extract_provenance.register(frozenset) @extract_provenance.register(set) def _extract_provenance_set(x): @@ -140,7 +161,7 @@ def _extract_provenance_set(x): @extract_provenance.register(tuple) @extract_provenance.register(dict) def _extract_provenance_pytree(x): - flat_args, spec = tree_flatten(x) + flat_args, spec = tree_flatten(x, is_leaf=_is_size) xs = [] provenance = frozenset() for x, p in map(extract_provenance, flat_args): diff --git a/tests/ops/test_provenance.py b/tests/ops/test_provenance.py index 478ae27f05..05e70646fb 100644 --- a/tests/ops/test_provenance.py +++ b/tests/ops/test_provenance.py @@ -4,7 +4,16 @@ import pytest import torch -from pyro.ops.provenance import ProvenanceTensor, get_provenance, track_provenance +import pyro +import pyro.distributions as dist +from pyro.infer.inspect import get_model_relations +from pyro.ops.provenance import ( + ProvenanceTensor, + detach_provenance, + extract_provenance, + get_provenance, + track_provenance, +) from tests.common import assert_equal, requires_cuda @@ -66,3 +75,53 @@ def test_track_provenance(x): old_provenance = get_provenance(x) provenance = old_provenance | new_provenance assert provenance == get_provenance(track_provenance(x, new_provenance)) + + +@pytest.mark.parametrize( + "x", + [ + torch.Size([3]), + torch.Size([2, 5]), + torch.Size([]), + [(torch.Size([3]),), {}], + [(), {"shape": torch.Size([2, 5])}], + [(torch.zeros(2), torch.Size([3])), {}], + ], + ids=["size", "size_2d", "size_empty", "in_args", "in_kwargs", "mixed_with_tensor"], +) +def test_provenance_preserves_torch_size(x): + """torch.Size subclasses tuple, so it must not be rebuilt as a plain tuple.""" + + def assert_sizes_intact(original, result): + if isinstance(original, torch.Tensor): + return # tensors are legitimately wrapped/unwrapped + assert type(original) is type(result) + if isinstance(original, torch.Size): + assert result == original + elif isinstance(original, (list, tuple)): + for a, b in zip(original, result): + assert_sizes_intact(a, b) + elif isinstance(original, dict): + for key in original: + assert_sizes_intact(original[key], result[key]) + + assert_sizes_intact(x, track_provenance(x, frozenset({"a"}))) + assert_sizes_intact(x, extract_provenance(x)[0]) + assert_sizes_intact(x, detach_provenance(x)) + + +def test_multinomial_render_model(): + """https://github.com/pyro-ppl/pyro/issues/3436""" + + def model(): + probs = pyro.param("probs", torch.tensor([0.1, 0.2, 0.7])) + return pyro.sample("y", dist.Multinomial(total_count=10, probs=probs)) + + relations = get_model_relations(model) + assert "y" in relations["sample_sample"] or "y" in relations["sample_param"] + + +def test_provenance_tensor_new_with_size(): + """Tensor.new() reads a torch.Size as a shape but a tuple as data.""" + x = ProvenanceTensor(torch.zeros(4, dtype=torch.long), frozenset({"x"})) + assert tuple(x.new(torch.Size([3])).shape) == (3,)