pyro-ppl / numpyro / 29599788232

92%

push

github

Miscellaneous bugfixes (#2215)

* infer.svi: treat empty JaxTrees of mutables as equivalent to None at all

Signed-off-by: Eli Sennesh <elisennesh@gmail.com>

* notebooks.source.nnx_example: employ TraceGraph_ELBO to hit mutables-check

Signed-off-by: Eli Sennesh <elisennesh@gmail.com>

* numpyro.primitives: correct implementation and use of callable init_value()

Signed-off-by: Eli Sennesh <elisennesh@gmail.com>

* numpyro.distributions: intermediates through log_prob of Expanded, Independent

TransformedDistribution.sample_with_intermediates caches the pre-transform
base sample so that log_prob(value, intermediates) can reuse it instead of
recomputing transform.inv(value). That optimization was silently lost when the
transformed distribution was wrapped:

- Independent inherited the base no-op sample_with_intermediates (returning an
  empty []), so the cached value never reached the trace, and its log_prob did
  not accept intermediates.
- ExpandedDistribution.log_prob ignored the intermediates it was passed (it
  carried a "TODO: utilize intermediates").

Both remain backward compatible (intermediates defaults to None).

For a TransformedDistribution wrapped by to_event()/expand(), this lets log_prob
reuse the sampled base value rather than inverting the transform, a potential
speedup in certain cases.  Include a test to make sure this actually works.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Eli Sennesh <elisennesh@astera.org>

* numpyro.primitives: include type annotation, fix formatting in param()

Signed-off-by: Eli Sennesh <elisennesh@astera.org>

---------

Signed-off-by: Eli Sennesh <elisennesh@gmail.com>
Signed-off-by: Eli Sennesh <elisennesh@astera.org>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

14 of 16 new or added lines in 2 files covered. (87.5%)

14611 of 15819 relevant lines covered (92.36%)

1.69 hits per line