[bug] Multinomial distribution raises RuntimeError within pyro.render_model

[markmcleod-mf]

Issue Description

The multinomial distribution when passed a pyro.param as probabilities and sampled within a call of pyro.render_model raises a runtime error.

This does not occur without the render_model, or if not using a parameter, or with other distributions.

Environment

• Ubuntu 24.04,
• python 3.12
• torch 2.7.0
• pyro 1.9.1

Code Snippet

import pyro
import torch
import pyro.distributions as dist

def model():
    x = pyro.param("x",torch.tensor([0.1,0.2,0.7]))
    # Multinomial causes an error when within render_model
    y=pyro.sample("y", dist.Multinomial(total_count=10, probs=x))
    # Other distributions do not
    # y=pyro.sample("y", dist.Dirichlet(concentration=x))
    return y


print(model())  # tensor([1., 2., 7.]) or similar

viz = pyro.render_model(model)  # Error when using Multinomial distribution

[yuanx749]

I tracked down the error, and found that during render_model, ProvenanceTensor is used in distribution, but the event_shape becomes tuple instead of torch.Size, which causes the problem.

pyro/pyro/ops/provenance.py

Lines 67 to 70 in 1bbbf38

	def __torch_function__(cls, func, types, args=(), kwargs=None):
	_args, _kwargs = detach_provenance([args, kwargs or {}])
	ret = func(*_args, **_kwargs)
	return track_provenance(ret, get_provenance([args, kwargs]))

I believe implementing dispatch for torch.Size in track_provenance and extract_provenance could solve the problem.