pyro-ppl / numpyro / 36754303847

93%

Build Type

push

github

Committed by web-flow

Commit Message

fix(distributions): type DistributionMeta.__call__ to return the instance type (#2293)

* fix(distributions): type distribution sampling as jax.Array

`DistributionMeta.__call__` was annotated `-> Any`, so every constructed distribution (`dist.Normal(0.0, 1.0)`) was inferred as `Any` by ty and the `-> Array` annotations on concrete `sample` methods never reached callers. Return the constructed instance type via a bound `TypeVar` instead.

With that in place, tighten the sampling surface in `distribution.py` (`sample`, `rsample`, `sample_with_intermediates`, `__call__`) from `ArrayLike` to `Array`, and fix the latent ty errors that `Any` had been masking.

`Delta.sample` returned the raw `v` (Python scalar or numpy array) when `sample_shape` was empty; it now always returns a `jax.Array`, matching every other distribution. The `Delta` carve-out in `test_dist_shape` is removed accordingly.

* fix(distributions): keep ArrayLike on base-class sampling per gh-2206

The base and wrapper distributions (Distribution, ExpandedDistribution, Independent, MaskedDistribution, TransformedDistribution, Delta) intentionally keep `ArrayLike` on `sample`/`rsample`/`sample_with_intermediates`/`__call__`, and `Delta.sample` returns the stored value untouched, so that numpy-backed or string-valued custom distributions pass through handlers (see gh-2206 and `test_handlers_with_string_valued_distribution`). Restore that surface and the `Delta` carve-out in `test_dist_shape`.

The `DistributionMeta.__call__` return type is what lets the concrete `-> Array` annotations on each distribution's `sample` reach callers; that change and the ty fixes it surfaced remain.

* refactor(distributions): drop comment on DistributionMeta.__call__

Coverage Stats

13 of 15 new or added lines in 5 files covered. (86.67%)

1 existing line in 1 file now uncovered.

15643 of 16873 relevant lines covered (92.71%)

1.7 hits per line