pyro-ppl / numpyro / 33474642619

93%

Build Type

push

github

Committed by web-flow

Commit Message

Ensure distribution moments/samples are jax arrays (#2206)

* Add parametrized test that distribution outputs are jax arrays

test_output_is_array checks that sample/log_prob/mean/variance/entropy
return jax.Array rather than python/numpy scalars, parametrized per
method so the per-method coverage is visible. Methods whose output type
tracks the input (e.g. Poisson.mean is self.rate, Beta.mean is plain
python arithmetic) are recorded in OUTPUT_TYPE_XFAIL and marked strict
xfail, so the dict must be updated if any of them later returns an array.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Narrow method return annotations from ArrayLike to Array

The distribution methods sample/log_prob/mean/variance/entropy return jax
arrays for every distribution exercised by test_output_is_array, so their
return type is narrowed from ArrayLike to Array. ArrayLike is the correct
type for an input (a method does accept python/numpy scalars), but on an
output it is an under-specified promise: it admits bool/int/float/complex
that the method never returns and that a consumer cannot index or reshape.

Input parameter annotations are left as ArrayLike. The implementation-
dependent tail recorded in OUTPUT_TYPE_XFAIL (e.g. Poisson.mean is
self.rate, Beta.mean is plain python arithmetic) keeps ArrayLike, since
those genuinely return a python/numpy scalar when given scalar params.

TransformedDistribution.sample keeps a `# type: ignore[return-value]`
rather than coercing, because Transform.__call__ is typed to return the
wide ArrayLike while the value is already a jax array at runtime.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Parametrize output-type test over jit and assert arrays under jit

Under jax.jit every method output is materialised as a jax array, so the
moment exceptions in OUTPUT_TYPE_XFAIL only apply in the non-jitted case;
the jitted case asserts every method returns a jax array with no exceptions.

Co-Authored-By: Claude Opus 4... (continued)

Coverage Stats

720 of 727 new or added lines in 14 files covered. (99.04%)

1 existing line in 1 file now uncovered.

14892 of 16098 relevant lines covered (92.51%)

1.68 hits per line