Skip to content

Apply init_to_median in unconstrained space.#1994

Merged
fehiepsi merged 4 commits intopyro-ppl:masterfrom
tillahoffmann:init-unconstrained
Mar 6, 2025
Merged

Apply init_to_median in unconstrained space.#1994
fehiepsi merged 4 commits intopyro-ppl:masterfrom
tillahoffmann:init-unconstrained

Conversation

@tillahoffmann
Copy link
Copy Markdown
Collaborator

@tillahoffmann tillahoffmann commented Mar 3, 2025

The elementwise median of a random variable does not necessarily belong to the support of the generating distribution. I recently ran into this issue with ZeroSumNormal distributions where init_to_median gave samples that did not sum to zero.

This PR updates init_to_median to work in the unconstrained space, guaranteeing valid samples. It also fixes a small bug where init_to_mean does not return a value.

This change should not change the initialization given fixed random key for simple bijections to the support, but it will affect more complex transformations when median(transform(x)) != transform(median(x)).

sample_shape=(num_samples,) + sample_shape, rng_key=rng_key
)
return jnp.median(samples, axis=0)
from numpyro.infer.util import helpful_support_errors
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just curious, you put the import here because of circular import?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I borrowed that line from init_to_uniform here.

# XXX: we import here to avoid circular import
from numpyro.infer.util import helpful_support_errors

@fehiepsi fehiepsi merged commit d393f5a into pyro-ppl:master Mar 6, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants