Preserve torch.Size through provenance tracking

[Mohit-Ak]

Fixes #3436.

The problem

torch.Size subclasses tuple. That means it matches the list/tuple/dict branches of track_provenance and extract_provenance, which run it through tree_flatten/tree_unflatten. Those decompose the Size into its plain int leaves and rebuild it as an ordinary tuple, so any torch.Size passing through the provenance machinery comes back out with its type erased:

>>> from pyro.ops.provenance import extract_provenance
>>> extract_provenance(torch.Size([3, 4]))[0]
(3, 4)          # a plain tuple, not torch.Size([3, 4])

That looks harmless, but torch overloads on exactly this distinction: Tensor.new() treats a torch.Size as a shape and a tuple as data.

>>> torch.zeros(3).new(torch.Size([3])).shape
torch.Size([3])   # shape
>>> torch.zeros(3).new((3,)).shape
torch.Size([1])   # 3 was read as data

Multinomial.sample() does counts = samples.new(self._extended_shape(sample_shape)).zero_(). Under render_model the sample is a ProvenanceTensor, so __torch_function__ routes the args through detach_provenance(), the shape degrades to a tuple, counts comes back with shape (1,) instead of (3,), and the next line fails:

RuntimeError: index 2 is out of bounds for dimension 0 with size 1
  File ".../torch/distributions/multinomial.py", line 123, in sample
    counts.scatter_add_(-1, samples, torch.ones_like(samples))
  File ".../pyro/ops/provenance.py", line 69, in __torch_function__
    ret = func(*_args, **_kwargs)

This is why the issue only shows up for Multinomial and only inside render_model — you need a distribution that reconstructs a tensor from its own shape, plus provenance tracking being active.

The fix

Two parts, and both are needed:

1. Register torch.Size on track_provenance and extract_provenance. A Size holds plain ints, so there is no provenance to add or extract — return it unchanged.
2. Pass is_leaf=_is_size to the tree_map/tree_flatten calls in the pytree branches.

Part 2 is the one that actually repairs the reported crash. The failing path is detach_provenance([args, kwargs]), where the Size is nested inside a list, so the pytree helper flattens it before singledispatch ever gets a chance to see it. I checked this rather than assuming — registering the dispatch alone leaves the bug in place for every nested case:

input	current	dispatch only	dispatch + is_leaf
bare Size	tuple	Size	Size
Size in args tuple	tuple	tuple	Size
Size nested in list	tuple	tuple	Size
Size in kwargs	tuple	tuple	Size
plain tuple (control)	tuple	tuple	tuple

The control row confirms ordinary tuples still flatten normally.

Testing

Added to tests/ops/test_provenance.py: a parametrized round-trip test covering bare/2-D/empty Size plus nesting in args, kwargs, and mixed with a tensor; a Tensor.new() test pinning the shape-vs-data behavior; and a get_model_relations test that is the issue's reproducer.

RED→GREEN, verified in a detached worktree so the buggy code ran against the new tests:

# pristine provenance.py + new tests
8 failed, 8 passed, 64 skipped        <- the 8 passing are the pre-existing control
# with the fix
16 passed, 64 skipped

Wider suites at the committed SHA:

tests/ops/test_provenance.py                    16 passed, 64 skipped
tests/infer/test_inspect.py                     25 passed
tests/infer/test_compute_downstream_costs.py    47 passed
tests/ops/                                      4607 passed, 3 failed, 76 skipped

The 3 failures are test_gaussian_funsor — ModuleNotFoundError: No module named 'funsor'. They fail identically on pristine upstream (checked in a clean worktree), so they're a missing optional dep, not this change.

make lint gates on the touched files: ruff check clean, black --check clean, mypy pyro/ops/provenance.py clean.