Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

Easier Use of Distributions Under `vmap`

Open
  • Is something wrong?
#2,271 3 comments 4 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

@DanWaxman is already working on this.

Since Sep 18, 2026.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
35/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
numpy, python
Domain
machine-learning

Research direction

Start by reviewing the distribution shape-handling code and tests referenced in the issue, especially test_distributions/test_promote_batch_shape_shares_data_and_preserves_input, alongside the dw-exploring-vmap-dists branch. Compare the vmap example with behavior around expanded and transformed distributions; done means the example works without breaking the cited shape and data-preservation semantics and the relevant tests are reliable.

Written by the indexing model from the issue text.

Description

enhancement

Distributions carry some static metadata (e.g., the batch shape), that make it a bit annoying to use in some other contexts. For example, they sometimes cause issues in dynestyx, where we use numpyro distributions quite generally.

An example of something I'd like to work, but doesn't, is just accessing variables of vmaped distributions:

def foo(mean):
    return dist.Normal(mean, jnp.ones_like(mean))

jax.vmap(foo)(jnp.zeros((2,2))).mean

Speaking offline with @fehiepsi, we agreed it would be desirable for such things to "just work." Looking back at previous discussions, however, it seems this is expected behavior under the current implementation (cf. https://github.com/pyro-ppl/numpyro/pull/1529#issuecomment-1411377292 or https://github.com/pyro-ppl/numpyro/issues/1684). Nevertheless, it would be nice to have a more extendable way of working with such things than imposing lots of numpyro-specific helpers on downstream users.

One way that might be possible is to try and defer shape evaluations when they may be inferred from parameter shapes (e.g., in all the scalar cases), in which case they become a lazy property. I had a quick vibe-coded go at this and it seems possible (see https://github.com/DanWaxman/numpyro/tree/dw-exploring-vmap-dists), though I haven't taken the time to go through that code very carefully (in particular, things that happen around expanded and transformed distributions), and some of the tests look wonky (in particular, I think Codex just changed test_distributions/test_promote_batch_shape_shares_data_and_preserves_input to pass ignoring all semantics haha), so please don't interpret it as a PR.

If that route seems interesting/acceptable, I'd be happy to take a closer/more careful look to see what it would take to do well and what it may break.

Dominant language
Python
Stars
2.8k
Forks
334
Avg merge
2d 15h
Merged PRs (30d)
71

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from pyro-ppl/numpyro

All issues in pyro-ppl/numpyro

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.