Sampling from MixtureSameFamily is slow when num_components is large.

Sampling from a mixture distribution with many components can be quite fast since you only need to sample one component from the categorical distribution and then sample from that component. However, in the implementation of MixtureSameFamily (in _sample_n) a sample is taken from all components and then later masked out such that only the sample from a specific component remains (or several components when n > 1). This seems quite inefficient both in terms of computation time and memory requirement.
In my specific situation, I'm trying to sample from a posterior predictive distribution of a Bayesian neural network created with MCMC sampling. This means my predictive distribution is a mixture of a really large number of components and when I'm trying to get more than about 100 samples at once, my (16GB of) RAM is filled and my computer pretty much freezes.

I've hacked together a really quick and dirty work around for my situation like so:

class MixtureSameFamilySampleFix(tfd.MixtureSameFamily):
    def _sample_n(self, n, seed):
        components_seed, mix_seed = samplers.split_seed(seed,
                                                        salt='MixtureSameFamily')
        try:
          seed_stream = SeedStream(seed, salt='MixtureSameFamily')
        except TypeError as e:  # Can happen for Tensor seeds.
          seed_stream = None
          seed_stream_err = e                                                
        try:
          mix_sample = self.mixture_distribution.sample(
              n, seed=mix_seed)  # [n, B] or [n]
        except TypeError as e:
          if ('Expected int for argument' not in str(e) and
              TENSOR_SEED_MSG_PREFIX not in str(e)):
            raise
          if seed_stream is None:
            raise seed_stream_err
          msg = ('Falling back to stateful sampling for `mixture_distribution` '
                 '{} of type `{}`. Please update to use `tf.random.stateless_*` '
                 'RNGs. This fallback may be removed after 20-Aug-2020. ({})')
          warnings.warn(msg.format(self.mixture_distribution.name,
                                   type(self.mixture_distribution),
                                   str(e)))
          mix_sample = self.mixture_distribution.sample(
              n, seed=seed_stream())  # [n, B] or [n]
        _seed = int(components_seed[0].numpy())
        ret = tf.stack([self.components_distribution[i_component.numpy()].sample(seed=_seed + i) for i, i_component in enumerate(mix_sample)], axis=0)
        return ret
# Only the last three lines are mine, everything else I've copied together from the existing code of MixtureSameFamily

In my current code this is two orders of magnitude faster and basically uses up no RAM, but of course this is a really bad piece of code (I have no idea what I'm doing with the random seeds for example and it uses a for loop).
I'd appreciate it if you made that more efficient on your end! :)

Start in MixtureSameFamily._sample_n, the implementation entry point named in the issue, and compare its current component sampling with the supplied workaround. The change is done when sampling selects and samples only the chosen components without materializing samples from every component, while preserving correct random-seed handling and reducing memory use for large mixtures.