pyro-ppl / numpyro / 27173226273

92%

push

github

Fix GaussianStateSpace.variance shape when batch comes from initial_value (#2205)

GaussianStateSpace.variance was computed from scale_tril alone, dropping any
batch dimension contributed by initial_value. When the batch shape originates
from initial_value (e.g. scale_tril (2,2) with initial_value (3,2)), variance
came out (num_steps, state_dim) instead of batch_shape + event_shape.

Broadcast the result to batch_shape + event_shape. The initial value is
deterministic and does not affect variance values, only the shape, so
broadcasting is exact.

Also add a uniform moment-shape contract check in test_mean_var asserting
mean/variance have shape batch_shape + event_shape for every distribution.

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

2 of 2 new or added lines in 1 file covered. (100.0%)

14409 of 15640 relevant lines covered (92.13%)

1.68 hits per line