Joint

class pmrf.distributions.Joint(distributions: PyTree[distreqx.distributions._distribution.AbstractDistribution])

Bases: AbstractDistribution

Joint distribution over a PyTree of statistically independent distributions.

Samples from the Joint distribution take the form of a PyTree structure that matches the structure of the underlying distributions. Log-probabilities, entropies, and KL divergences are summed over the tree.

Initializes a Joint distribution.

Arguments:

distributions: A PyTree of distreqx distributions.

cdf(value: PyTree[jax.jaxlib._jax.Array]) → Array: Evaluates the joint cumulative distribution function.

entropy() → Array: Calculates the sum of Shannon entropies (in nats).

icdf(value: PyTree[jax.jaxlib._jax.Array]) → PyTree[jax.jaxlib._jax.Array]: Evaluates the joint inverse cumulative distribution function.

kl_divergence(other_dist, **kwargs) → Array: Calculates the KL divergence between two Joint distributions.

log_cdf(value: PyTree[jax.jaxlib._jax.Array]) → Array: Evaluates the log of the joint CDF.

log_prob(value: PyTree[jax.jaxlib._jax.Array]) → Array: Compute the total log probability of the distributions in the tree.

log_survival_function(value: PyTree[jax.jaxlib._jax.Array]) → Array: Evaluates the log of the joint survival function.

mean() → PyTree[jax.jaxlib._jax.Array]: Calculates the joint mean.

median() → PyTree[jax.jaxlib._jax.Array]: Calculates the joint median.

mode() → PyTree[jax.jaxlib._jax.Array]: Calculates the joint mode.

prob(value: PyTree[jax.jaxlib._jax.Array]) → Array: Calculates the total probability of a joint event.

sample(key: Key[jaxlib._jax.Array, '']) → PyTree[jax.jaxlib._jax.Array]: Samples a joint event.

sample_and_log_prob(key: Key[jaxlib._jax.Array, '']) → tuple[PyTree[jax.jaxlib._jax.Array], Array]: Returns a joint sample and its total log prob.

stddev() → PyTree[jax.jaxlib._jax.Array]: Calculates the joint standard deviation.

survival_function(value: PyTree[jax.jaxlib._jax.Array]) → Array: Evaluates the joint survival function.

variance() → PyTree[jax.jaxlib._jax.Array]: Calculates the joint variance.

distributions: PyTree[distreqx.distributions._distribution.AbstractDistribution]

property event_shape: PyTree[tuple]: Shape of the joint event.

property support: tuple[PyTree[jax.jaxlib._jax.Array], PyTree[jax.jaxlib._jax.Array]]: See Distribution.support.