Table of Contents
Shortcuts

# Source code for torch.distributions.mixture_same_family

    from typing import Dict

    import torch
    from torch.distributions import Categorical, constraints
    from torch.distributions.distribution import Distribution

    __all__ = ["MixtureSameFamily"]

    class MixtureSameFamily(Distribution):
        r"""
        The `MixtureSameFamily` distribution implements a (batch of) mixture
        distribution where all component are from different parameterizations of
        the same distribution type. It is parameterized by a `Categorical`
        "selecting distribution" (over `k` component) and a component
        distribution, i.e., a `Distribution` with a rightmost batch shape
        (equal to `[k]`) which indexes each (batch of) component.

        Examples::

            >>> # xdoctest: +SKIP("undefined vars")
            >>> # Construct Gaussian Mixture Model in 1D consisting of 5 equally
            >>> # weighted normal distributions
            >>> mix = D.Categorical(torch.ones(5,))
            >>> comp = D.Normal(torch.randn(5,), torch.rand(5,))
            >>> gmm = MixtureSameFamily(mix, comp)

            >>> # Construct Gaussian Mixture Model in 2D consisting of 5 equally
            >>> # weighted bivariate normal distributions
            >>> mix = D.Categorical(torch.ones(5,))
            >>> comp = D.Independent(D.Normal(
            ...          torch.randn(5,2), torch.rand(5,2)), 1)
            >>> gmm = MixtureSameFamily(mix, comp)

            >>> # Construct a batch of 3 Gaussian Mixture Models in 2D each
            >>> # consisting of 5 random weighted bivariate normal distributions
            >>> mix = D.Categorical(torch.rand(3,5))
            >>> comp = D.Independent(D.Normal(
            ...         torch.randn(3,5,2), torch.rand(3,5,2)), 1)
            >>> gmm = MixtureSameFamily(mix, comp)

        Args:
            mixture_distribution: `torch.distributions.Categorical`-like
                instance. Manages the probability of selecting component.
                The number of categories must match the rightmost batch
                dimension of the `component_distribution`. Must have either
                scalar `batch_shape` or `batch_shape` matching
                `component_distribution.batch_shape[:-1]`
            component_distribution: `torch.distributions.Distribution`-like
                instance. Right-most batch dimension indexes component.
        """
        arg_constraints: Dict[str, constraints.Constraint] = {}
        has_rsample = False

        def __init__(
            self, mixture_distribution, component_distribution, validate_args=None
        ):
            self._mixture_distribution = mixture_distribution
            self._component_distribution = component_distribution

            if not isinstance(self._mixture_distribution, Categorical):
                raise ValueError(
                    " The Mixture distribution needs to be an "
                    " instance of torch.distributions.Categorical"
                )

            if not isinstance(self._component_distribution, Distribution):
                raise ValueError(
                    "The Component distribution need to be an "
                    "instance of torch.distributions.Distribution"
                )

            # Check that batch size matches
            mdbs = self._mixture_distribution.batch_shape
            cdbs = self._component_distribution.batch_shape[:-1]
            for size1, size2 in zip(reversed(mdbs), reversed(cdbs)):
                if size1 != 1 and size2 != 1 and size1 != size2:
                    raise ValueError(
                        f"`mixture_distribution.batch_shape` ({mdbs}) is not "
                        "compatible with `component_distribution."
                        f"batch_shape`({cdbs})"
                    )

            # Check that the number of mixture component matches
            km = self._mixture_distribution.logits.shape[-1]
            kc = self._component_distribution.batch_shape[-1]
            if km is not None and kc is not None and km != kc:
                raise ValueError(
                    f"`mixture_distribution component` ({km}) does not"
                    " equal `component_distribution.batch_shape[-1]`"
                    f" ({kc})"
                )
            self._num_component = km

            event_shape = self._component_distribution.event_shape
            self._event_ndims = len(event_shape)
            super().__init__(
                batch_shape=cdbs, event_shape=event_shape, validate_args=validate_args
            )

        def expand(self, batch_shape, _instance=None):
            batch_shape = torch.Size(batch_shape)
            batch_shape_comp = batch_shape + (self._num_component,)
            new = self._get_checked_instance(MixtureSameFamily, _instance)
            new._component_distribution = self._component_distribution.expand(
                batch_shape_comp
            )
            new._mixture_distribution = self._mixture_distribution.expand(batch_shape)
            new._num_component = self._num_component
            new._event_ndims = self._event_ndims
            event_shape = new._component_distribution.event_shape
            super(MixtureSameFamily, new).__init__(
                batch_shape=batch_shape, event_shape=event_shape, validate_args=False
            )
            new._validate_args = self._validate_args
            return new

        @constraints.dependent_property
        def support(self):
            # FIXME this may have the wrong shape when support contains batched
            # parameters
            return self._component_distribution.support

        @property
        def mixture_distribution(self):
            return self._mixture_distribution

        @property
        def component_distribution(self):
            return self._component_distribution

        @property
        def mean(self):
            probs = self._pad_mixture_dimensions(self.mixture_distribution.probs)
            return torch.sum(
                probs * self.component_distribution.mean, dim=-1 - self._event_ndims
            )  # [B, E]

        @property
        def variance(self):
            # Law of total variance: Var(Y) = E[Var(Y|X)] + Var(E[Y|X])
            probs = self._pad_mixture_dimensions(self.mixture_distribution.probs)
            mean_cond_var = torch.sum(
                probs * self.component_distribution.variance, dim=-1 - self._event_ndims
            )
            var_cond_mean = torch.sum(
                probs * (self.component_distribution.mean - self._pad(self.mean)).pow(2.0),
                dim=-1 - self._event_ndims,
            )
            return mean_cond_var + var_cond_mean

        def cdf(self, x):
            x = self._pad(x)
            cdf_x = self.component_distribution.cdf(x)
            mix_prob = self.mixture_distribution.probs

            return torch.sum(cdf_x * mix_prob, dim=-1)

        def log_prob(self, x):
            if self._validate_args:
                self._validate_sample(x)
            x = self._pad(x)
            log_prob_x = self.component_distribution.log_prob(x)  # [S, B, k]
            log_mix_prob = torch.log_softmax(
                self.mixture_distribution.logits, dim=-1
            )  # [B, k]
            return torch.logsumexp(log_prob_x + log_mix_prob, dim=-1)  # [S, B]

        def sample(self, sample_shape=torch.Size()):
            with torch.no_grad():
                sample_len = len(sample_shape)
                batch_len = len(self.batch_shape)
                gather_dim = sample_len + batch_len
                es = self.event_shape

                # mixture samples [n, B]
                mix_sample = self.mixture_distribution.sample(sample_shape)
                mix_shape = mix_sample.shape

                # component samples [n, B, k, E]
                comp_samples = self.component_distribution.sample(sample_shape)

                # Gather along the k dimension
                mix_sample_r = mix_sample.reshape(
                    mix_shape + torch.Size([1] * (len(es) + 1))
                )
                mix_sample_r = mix_sample_r.repeat(
                    torch.Size([1] * len(mix_shape)) + torch.Size([1]) + es
                )

                samples = torch.gather(comp_samples, gather_dim, mix_sample_r)
                return samples.squeeze(gather_dim)

        def _pad(self, x):
            return x.unsqueeze(-1 - self._event_ndims)

        def _pad_mixture_dimensions(self, x):
            dist_batch_ndims = self.batch_shape.numel()
            cat_batch_ndims = self.mixture_distribution.batch_shape.numel()
            pad_ndims = 0 if cat_batch_ndims == 1 else dist_batch_ndims - cat_batch_ndims
            xs = x.shape
            x = x.reshape(
                xs[:-1]
                + torch.Size(pad_ndims * [1])
                + xs[-1:]
                + torch.Size(self._event_ndims * [1])
            )
            return x

        def __repr__(self):
            args_string = (
                f"\n  {self.mixture_distribution},\n  {self.component_distribution}"
            )
            return "MixtureSameFamily" + "(" + args_string + ")"

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.