Ensure distribution moments/samples are jax arrays - #2206

Open

tillahoffmann wants to merge 5 commits into

pyro-ppl:masterfrom

tillahoffmann:add-jax-array-output-test

Commits on Jul 2, 2026

Add parametrized test that distribution outputs are jax arrays

tillahoffmann
and
claude
committed
ad5ddd6
Narrow method return annotations from ArrayLike to Array

tillahoffmann
and
claude
committed
92b51d0
Parametrize output-type test over jit and assert arrays under jit

tillahoffmann
and
claude
committed
1435475
Promote mean and variance to jax.Array and narrow annotations.
tillahoffmann
committed
98795d6
Return ArrayLike from dual-backend Constraint/Transform methods

tillahoffmann
and
claude
committed
8f09429