Adapting GMM example to higher dimensions

  • What tutorial are you running?
    Gaussian Mixture Model
    Gaussian Mixture Model — Pyro Tutorials 1.8.4 documentation
  • What version of Pyro are you using?
    dev branch
  • Please link or paste relevant code, and steps to reproduce.
    The GMM example shows a 1D density estimation example, so naturally I wanted to try to modify it to have additional dimensions. I have tried making various modifications to the model, but with no luck so far.
    For all of these, I modified the data to a [5x2] array → 5 data points in 2d

Ver 1: simply making the prior for the locs to 2D independent Gaussian

def model(data):
    # Global variables.
    weights = pyro.sample('weights', dist.Dirichlet(0.5 * torch.ones(K)))
    scale = pyro.sample('scale', dist.LogNormal(0., 2.))
    with pyro.iarange('components', K):
        locs = pyro.sample('locs', dist.Normal(torch.zeros(2), 10.))
    with pyro.iarange('data', len(data)):
        # Local variables.
        assignment = pyro.sample('assignment', dist.Categorical(weights))
        pyro.sample('obs', dist.Normal(locs[assignment], scale), obs=data)

results in error complaining about broadcast:
line 33, in _pyro_sample: f.name, msg['name'], f.dim, f.size, target_batch_shape[f.dim]))ErrorCode: ValueError: Shape mismatch inside plate('data') at site obs dim -1, 5 vs 2

Ver 2: specify dimension of the ‘data’ plate to be on the left most dimension (aka -2)

 def model(data):
    # Global variables.
    weights = pyro.sample('weights', dist.Dirichlet(0.5 * torch.ones(K)))
    scale = pyro.sample('scale', dist.LogNormal(0., 2.))
    with pyro.iarange('components', K):
        locs = pyro.sample('locs', dist.Normal(torch.zeros(2), 10.))

    with pyro.iarange('data', len(data), dim=-2):
        # Local variables.
        assignment = pyro.sample('assignment', dist.Categorical(weights))
        pyro.sample('obs', dist.Normal(locs[assignment], scale), obs=data)

results in ValueError: Model and guide shapes disagree at site 'locs': torch.Size([2]) vs torch.Size([2, 2]) because the autoguide for some reason don’t like the higher dim locs?

Ver 2.1: add “dim=-2” to the components since that is supposed to be the independent dimension

def model(data):
    # Global variables.
    weights = pyro.sample('weights', dist.Dirichlet(0.5 * torch.ones(K)))
    scale = pyro.sample('scale', dist.LogNormal(0., 2.))
    with pyro.iarange('components', K, dim=-2):
        locs = pyro.sample('locs', dist.Normal(torch.zeros(2), 10.))

    with pyro.iarange('data', len(data), dim=-2):
        # Local variables.
        assignment = pyro.sample('assignment', dist.Categorical(weights))
        pyro.sample('obs', dist.Normal(locs[assignment], scale), obs=data)

results in

ValueError: Error while packing tensors at site 'locs':
  Invalid tensor shape.
  Allowed dims: -2
  Actual shape: (2, 2)

I have also tried number of other options such as moving the locs outside the plate (it then complains about the assignments), moving assignments into its own plate with dim=-1 (since the assignments are 1-D while obs is 2-D)…basically nothing worked. Can someone point to a working example or give me a direction as to how this is supposed to work?

Hi @jeanlancel, you should be able to add an .independent() dimension along the lines of

def model(data):
    # Global variables.
    weights = pyro.sample('weights', dist.Dirichlet(0.5 * torch.ones(K)))
    scale = pyro.sample('scale', dist.LogNormal(0., 2.))
    with pyro.plate('components', K):
        locs = pyro.sample('locs', dist.Normal(torch.zeros(2), 10.).independent(1))
        assert locs.shape == (K, 2)
    with pyro.plate('data', len(data)):
        # Local variables.
        assignment = pyro.sample('assignment', dist.Categorical(weights))
        pyro.sample('obs', dist.Normal(locs[assignment], scale).independent(1),
                    obs=data)

(Note I’ve using the newer pyro.plate syntax rather than the older equivalent pyro.iarange. We’re planning to release Pyro 0.3 next week based on PyTorch 1.0.)