Add Deterministic and Joint#70
Conversation
gvcallen
changed the title
Deterministic places all probability mass on a single point (with an atol/rtol tolerance window for closeness comparisons). Joint composes a pytree of statistically independent distributions, summing log-probabilities/entropies/KL divergences over the tree. Ported from distrax as reference.
gvcallen
force-pushed
the
deterministic_and_joint
branch
from
c177f9a to
c88707d
Compare
| point (the `loc`). | ||
| """ | ||
|
|
||
| loc: Array |
| atol: Array | ||
| rtol: Array | ||
|
|
||
| def __init__(self, loc: Array | float, atol: float = 0.0, rtol: float = 0.0): |
There was a problem hiding this comment.
can we do with fields or would look bad here?
| - `distributions`: A PyTree of `distreqx` distributions. | ||
| """ | ||
| # Ensure there is at least one distribution in the tree | ||
| leaves = jax.tree_util.tree_leaves(distributions, is_leaf=_is_dist) |
| lambda dist: dist.entropy(), self.distributions, is_leaf=_is_dist | ||
| ) | ||
| entropies_summed = jax.tree_util.tree_map(jnp.sum, entropies) | ||
| return jnp.sum(jnp.asarray(jax.tree_util.tree_leaves(entropies_summed))) |
There was a problem hiding this comment.
this seems similar to independent which also does this reduction over leaves? return _reduce_helper(self.distribution.entropy())
what is the key distinction here between them?
There was a problem hiding this comment.
They are very similar, but I think the main difference is:
- Independent treats the underlying distributions as homogeneous (e.g. all Normal) and transforms the underlying distributions as to treat the resultant event as a single concept (for example, log prob returns a single number)
- Joint simply "groups" multiple heterogenous distributions/events (e.g. Normal, Uniform etc.) together (so you get multiple log probs representing that same event shape e.g. a pytree)
For general PyTree grouping of distributions, Joint is used (the exact converse to the Leafwise bijector) whereas to create e.g. multidimensional, independent normals you would use Independent (a distribution-specific concept).
There was a problem hiding this comment.
ah this makes sense, joint is over different distributions
There was a problem hiding this comment.
should document this more in the class docstring
| lambda dist: dist.entropy(), self.distributions, is_leaf=_is_dist | ||
| ) | ||
| entropies_summed = jax.tree_util.tree_map(jnp.sum, entropies) | ||
| return jnp.sum(jnp.asarray(jax.tree_util.tree_leaves(entropies_summed))) |
There was a problem hiding this comment.
ah this makes sense, joint is over different distributions
| lambda dist: dist.entropy(), self.distributions, is_leaf=_is_dist | ||
| ) | ||
| entropies_summed = jax.tree_util.tree_map(jnp.sum, entropies) | ||
| return jnp.sum(jnp.asarray(jax.tree_util.tree_leaves(entropies_summed))) |
There was a problem hiding this comment.
should document this more in the class docstring
| self.atol = jnp.asarray(atol) | ||
| self.rtol = jnp.asarray(rtol) | ||
|
|
||
| if self.loc.ndim != 0: |
There was a problem hiding this comment.
In principle this could be non scalar? we can refine the type hints/docstring to clarify either way
2 participants
Adds a Deterministic (point mass) and Joint distribution over pytrees. Reference code is from distrax (Deterministic and Joint)