See this notebook for a runnable version: plate-vs-not.ipynb · GitHub
def model1(data: Optional["np.ndarray[float]"] = None) -> None:
mus = ny.sample("mus", dist.Normal(0, 1).expand((2,)))
scale_tril = ny.sample("scale_tril", dist.LKJCholesky(2, concentration=1))
xs = ny.sample(
"xs", dist.MultivariateNormal(loc=mus, scale_tril=scale_tril), obs=data
)
print(xs.shape)
def model2(data: "np.ndarray[float]") -> None:
mus = ny.sample("mus", dist.Normal(0, 1).expand((2,)))
scale_tril = ny.sample("scale_tril", dist.LKJCholesky(2, concentration=1))
with ny.plate("obs", data.shape[0], dim=-1):
xs = ny.sample(
"xs", dist.MultivariateNormal(loc=mus, scale_tril=scale_tril), obs=data
)
print(xs.shape)
What is the semantic difference between observing (M, N)-shaped data inside an M-sized plate with an N-event-shaped distribution (model 2) and observing (M, N)-shaped data with an N-event-shaped distribution and no plate (model 1)? It seems like in the model 1 case, the sample is automatically getting broadcast to the obs
shape somehow? Is it automatically getting plated? As the notebook shows, at least for this scenario, the two approaches produce the same inference results.