Fix multidimensional Gaussian sample shapes

[AHMETHAKANBEZIR1]

Description

GaussianDistribution.sample(key, sample_shape=(2, 3)) raises a matrix-vector shape error. With (2, 2) and a non-diagonal covariance, it instead applies the covariance root along a sample axis and returns incorrect values. The single vmap only removes one sample axis before passing its input to the linear operator.

Flatten the sample axes before the affine transform, then restore the original noise shape. The scalar and one-dimensional sampling paths retain their behavior. Empty sample axes and zero-dimensional events also return the requested shape. Add a release note.

The regression compares eager and JIT samples with loc + normal(key, shape) @ cholesky(covariance).T, using the same key and a correlated covariance. It covers scalar, one-, two- and three-dimensional shapes, empty axes, and float32/float64 covariance parameters. On the base commit, 10 of 14 cases fail; 4 pass.

Related historical report: #236 (closed as stale, not fixed by that closure).

Validation

• Python 3.13.15, JAX 0.11.2, CPU; existing development environment through uv run --no-sync.
• Final distribution and KL tests: 54 passed, including four empty-event shape regressions.
• poe format and poe lint: passed.
• poe docstrings: 41 passed (one warning).
• Final full poe test with PYTHONUTF8=1: 3,247 passed, 1 skipped. An earlier run had four Windows cp1254 file-decoding failures; the UTF-8 setting resolves them without source changes.
• A separate coverage-instrumented invocation terminated without diagnostics locally; no coverage percentage is claimed.

GPU and the full documentation site were not run locally.

AI disclosure: this change and its tests were prepared and checked autonomously with OpenAI Codex on behalf of AHMETHAKANBEZIR1. No human code-review claim is made.