Combine irange + iarange + broadcast into pyro.plate

[fritzo]

New users of Pyro have often found it difficult to inderstand the interplay between iarange, irange, and @broadcast. This issue proposes to merge these three concepts into a single plate.

Historically @fritzo split map_data into irange+iarange, @neerajprad added @broadcast. Now that our usage patterns have converged, @eb8680 suggested making iarange and irange into messengers, in which case all can be fused.

Tasks

• Make iarange and irange into Messengers #1307 Make irange and iarange into messengers
• Fuse broadcast+iarange -> plate #1464 Fuse iarange + broadcast -> plate
• Replace irange by plate in docs, examples, tests #1550 Replace irange with plate (requires disambiguating lots of docs & error messages)
• Make plate into a handler? (requires replacing :class: with :func: in docstrings)

[fritzo]

@jwvdm I hope this interface is more intuitive!

[fritzo]

@ndgAtUber Are you ok with this renaming? We will maintain full backwards-compatibility.

[ngoodman]

could you point me to some examples of models written with this new interface?

(i like the cleanup and merging idea, but i want to think through the interface, to make sure it will be as clear as possible.)

[fritzo]

@ngoodman Take a look at test_valid_models.py.

[jwvdm]

@fritzo: I really like the name plate which is very intuitive, and the examples in test_valid_models.py look really clear!

When @esennesh and I went through your Tensor Shape Tutorial, the thing that was arguably the hardest to grok is the Distribution.independent construct. I personally struggled with this line:

with pyro.iarange("d_iarange", 3):
        d = pyro.sample("d", Normal(torch.zeros(3, 4, 5), 1).independent(2))

Part of this is due to the naming -- if I understand correctly independent(n) means "the last n dimensions are part of the event shape". In the above example, I would expect the syntax to be .independent(1) since the first dimension is part of the batch_shape whereas the remaining two dimensions are part of the event_shape.

The other thing that is just hard to wrap your head around is how plate aligns with dimensions in the batch shape. I actually think that you guys got this right, it just takes some getting used to.

Just to check my own comprehension here: What happens when you do

with pyro.iarange("d_iarange", 2, dim=-2):
        d = pyro.sample("d", Normal(torch.zeros(3, 4, 5), 1).independent(2))

This should yield batch_shape=(2, 3) and event_shape=(4, 5) is that correct?

A final note: the plate construct reminds me a lot of the Infer.NET ForEach construct. Here's the Infer.NET example for a GMM:

using (Variable.ForEach(n)) {  
  z[n] = Variable.Discrete(weights);  
  using (Variable.Switch(z[n])) {  
    data[n] = Variable.VectorGaussianFromMeanAndPrecision(  
      means[z[n]], precs[z[n]]);  
  }  
}

[fritzo]

@jwvdm Thanks for the feedback and pointer to Infer.NET ForEach! Your shapes are correct. I agree .independent() is confusing; we were following Tensorflow's Independent distribution but @jpchen is planning to rename it to something more intuitive #1406 before our next release. Suggestions welcome!

[jpchen]

I would expect the syntax to be .independent(1) since the first dimension is part of the batch_shape whereas the remaining two dimensions are part of the event_shape

you're not the first person to make this comment. i agree the original semantics would be more intuitive if independent specified the independent dimensions instead of the dependent ones. but since it was originally written that way it would be breaking to change the semantics as such, so we're probably going to just rename independent to be clear what it is doing.

[ngoodman]

i really like the move toward simplifying the interface!

i don't love the current idiom (though it's better than iarrange, etc). the nice thing about mapData and ForEach is that they feel very intuitive because they are explicit about what variable they bind.

for example in the simplest case:

with pyro.plate("data", data.shape[0]):
            assignments = pyro.sample("assignments", dist.Categorical(mix_proportions))
            pyro.sample("obs", dist.Normal(cluster_means[assignments], 1.), obs=data)

it feels like the with should be putting a datum into scope, which can then be used in the expected way -- if the plate was really a mapping over data, all the plurals should be singular (assignment, datum). didn't we once have the idiom something more like with pyro.plate(..) as datum?

i guess there is a bit of mismatch between the current sleekness where pyro.plate is about indexing and independence not specific values and my intuitions about iterating.... but maybe there is a happy middle ground? or a layer of helper functions that makes the standard iterating-over-data case clear?

[eb8680]

didn't we once have the idiom something more like with pyro.plate(..) as datum?

This behavior is still preserved, and you can still iterate over a plate sequentially just as with irange: for i in pyro.plate(...):; we've effectively just given iarange and irange the same name. The bigger change proposed in this issue (and added in #1307 and #1464 ) is broadcast-by-default.