Question on shape of `sample` method with added context

[phinate]

I've trained a conditional MAF with two additional context variables, pretty much identical to the two-moon example in the repo. However, I'm having some trouble coercing that context into the sample method.

The docstring for Distribution.sample is below, and adds that the samples will have output of shape with leading dimension determined by the context: https://github.com/bayesiains/nflows/blob/ac0bf432fc7904458a933ed14180f0ac26e3f93d/nflows/distributions/base.py#L54-L56

I'm not clear on why this is the case (I would not expect conditioning on side information to influence the shape of my samples), and I wasn't able to generate correct samples by supplying my two context variables either -- Distribution.sample seems to treat a context Tensor like [[y1, y2]] unusually, as I was getting bi-modal samples from a unimodal learned likelihood, almost like it was conditioning on [[y1, y1]] and [[y2 y2]] separately, and generating samples for each.

Apologies for no example (private data), but I can do any tests, though I suspect my error may be conceptual. Thanks for your time!

[phinate]

Uploading some images for comparison -- this is a bit of a spanner in the works for our current project, so would really appreciate any pointers!

[imurray]

The context should have shape [context_size, context_features], containing the context for context_size different examples, where you specified context_features when creating the flow.

sample(num_samples, context) then generates num_samples samples for each context. So if you provide 1000 contexts and ask for three samples, you'll get 3000 samples, in an array of shape [context_size, num_samples, features] = [1000, 3, 2]. It sounds like you want to provide one context, which you can do. You can then access the multiple samples for a particular context. If you want to lump all the samples together then do a .reshape([-1, 2]). It works for me.

Apologies for no example (private data)

It would be possible to include code for a toy example though.

[phinate]

Sagely advice @imurray, and sorry for not doing that sooner.

I modified examples/conditional_moons.ipynb to have a two dimensional context feature, then worked it through and things worked as your message describes. I then cross-referenced that line-by-line with our implementation, only to find an erroneous .reshape call that was flipping dimensions unecessarily! The only reason we got the plots above was purely coincidental when the shapes matched -- the samples are indeed not correct.

Thanks for your explanation though, I think your sentence

sample(num_samples, context) then generates num_samples samples for each context. So if you provide 1000 contexts and ask for three samples, you'll get 3000 samples, in an array of shape [context_size, num_samples, features] = [1000, 3, 2].

was what made the correctness of the shapes click for me. I would love to see that in the docstring if possible! I made the PR in #61.

[arturbekasov]

61 was merged -- resolving.