Sample from an infinite Gaussian Mixture model

ZSheikhb: <p>I have written a code which samples from the infinite Gaussian mixture model, using stick-breaking Dirichlet process. Here <code>z_sample_list</code> is a list of <code>K</code> components coming from the reparameterization trick for each component of <code>z</code><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/d/f/dfc2491cb88306bce81bf2f54d91121d1bcd6e7b.png" alt="Screen Shot 2022-03-24 at 4.50.12 PM" data-base62-sha1="vVsJjum8qoJ2YdKw2eAoHCgIcjh" width="432" height="90"><br>
The infinite Gaussian Mixture value is</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/0/001190ec029e7b2ee731796baee3c671b24570fc.png" data-download-href="https://discuss.pytorch.org/uploads/default/001190ec029e7b2ee731796baee3c671b24570fc" title="Screen Shot 2022-03-24 at 4.51.31 PM"><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/0/001190ec029e7b2ee731796baee3c671b24570fc.png" alt="Screen Shot 2022-03-24 at 4.51.31 PM" data-base62-sha1="BDnHxIz30o2eQT9H0zzfXpRww" width="690" height="73" data-dominant-color="F2F3F3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2022-03-24 at 4.51.31 PM</span><span class="informations">1140×122 7.69 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I think there is a problem in my implementation because the samples don’t look right. I will appreciate if somebody suggests how I can fix this code. I want my code to be able to sample from a distribution similar to <a href="https://github.com/pytorch/pytorch/blob/master/torch/distributions/mixture_same_family.py" rel="noopener nofollow ugc"><code>MixtureSameFamily</code></a> but use a similar method as reparameterization trick. Thanks.</p>
<pre><code class="lang-auto">def gather_nd(params, indices):
    ndim = indices.shape[-1]
    output_shape = list(indices.shape[:-1]) + list(params.shape[indices.shape[-1]:])
    flatted_indices = indices.view(-1, ndim)
    slices = [flatted_indices[:, i] for i in range(ndim)]
    slices += [Ellipsis]
    return params[slices].view(*output_shape)
</code></pre>
<pre><code>def mix_weights( beta):
    beta1m_cumprod = (1 - beta).cumprod(-1)
    return torch.mul(F.pad(beta, (0, 1), value=1) , F.pad(beta1m_cumprod, (1, 0), value=1))


def get_component_samples( batchSize,z_sample_list, z_dim, K, a, b, const_prior,device):
    
    # ensure stick segments sum to 1
    m = Beta(torch.ones_like(a, device=device, requires_grad=False), torch.ones_like(b, device=self.device, requires_grad=False)*const_prior])
    v_means =  m.sample().to(device=device)
    # compose into stick segments using pi = v \prod (1-v)
    pi_samples = mix_weights(v_means)[:,:-1]
    
    # sample a component index   
    
    component = torch.argmax( pi_samples  , 1).to(device=device, dtype=torch.int64)

    component = torch.cat( [torch.arange(0, batchSize, device=device).unsqueeze(1), component.unsqueeze(1)], 1)

    all_z = []
    for d in range(z_dim):
        temp_z = torch.cat( [z_sample_list[k][:, d].unsqueeze(1) for k in range(K)], dim=1)
        all_z.append(gather_nd(temp_z, component).unsqueeze(1))
    
    pi_samples         = pi_samples.unsqueeze(-1)
    pi_samples         = F.pad(input=pi_samples, pad=(0, 0, 1, 0), mode='constant', value=0)
    pi_samples         = pi_samples.expand(-1, K, 1)
    pi_samples         = pi_samples.permute(0,2,1)
    out      = torch.stack( all_z).to(device=device)
    out      = out.permute(1,0,2)
    concatenated_latent_space = torch.bmm(out, pi_samples)

    return torch.squeeze(torch.mean(concatenated_latent_space, 2, True))
</code></pre>