adding RandomVariable container class around Distribution

[ecotner]

• simple container class that encapsulates a Distribution
• arithmetic operators like +, -, /, **, etc. overloaded to perform transformations on the domain of the distribution
• easy method self.transform to apply other transformations and return another RandomVariable

Documented

• docstring for .rv that points to RandomVariable class
• add a file docs/source/contrib.randomvariable.rst following one of the other contrib.*.rst files
• add a contrib.randomvariable to docs/source/index.rst
• distring for .dist
• example usage with method chaining in RandomVariable

Tested

• distribution support tests

[eb8680]

@ecotner thanks for the PR. I stand by what I said in #2441, which is that this design breaks down when users (quite reasonably) try to build expressions that are not invertible or contain more than one random variable, especially if the random variables are not jointly independent. We could raise NotImplementedErrors or ValueErrors in these cases, but it's generally best to avoid designing expressive APIs that frequently violate user intuition in non-obvious ways.

Even disregarding the multiple-variables issue, I can also foresee a lot of feature creep and maintenance burden, up to reimplementing large swaths of the PyTorch tensor API as torch.distributions.Transforms. I seem to recall that this problem has been discussed in PyTorch issues about adding more complicated shape manipulations to torch.distributions.transforms, but I'm having trouble finding the issues at the moment. (edit: This issue in PyTorch about a core invertible op feature is also relevant: pytorch/pytorch#23756)

I think the best way to proceed with a feature that most closely resembles this PR is to resurrect the RandomVariable implementations in Funsor (pyro-ppl/funsor#130). There's been a lot of work done on both Funsor and PyTorch since then that would make wrapping distributions and ops/transforms generically (pyro-ppl/funsor#159) fairly straightforward, and I am happy to guide you through the process if you're interested.

However, I still think a better long-term solution to the multiple-variables problem is having pyro.sample statements emit funsor.Variables.

[fritzo]

@ecotner This is indeed a slick little DSL for specifying transforms, and I'd like to be able to expose it in Pyro. I think the safest way to do so in light of @eb8680 valid concerns is to:

1. Implement your DSL in a contrib package, say pyro.contrib.transform or pyro.contrib.randomvariable.
2. Clearly separate the Distribution DSL from the RandomVariable DSL, so as to avoid API creep and potential conflict with inference-oriented operations on distributions.

Regarding (2), we could see the Distribution DSL as providing chainable methods .to_event(), .mask(), and .expand(); whereas the RandomVariable DSL provides chainable methods .transform() .__add__() (for which we could add a torch-like alias .add()), .__mul__() etc. I think it would be pretty safe to add a single distribution propert .rv that switched from the Distribution DSL to the RandomVariable DSL, and provide a random variable property .dist to convert back to the Distribution DSL. Example usage could then be like

d = dist.Uniform(0, 1).rv.add(1).mul(0.5).tanh().dist.mask(my_mask).to_event(1)
#   \________________/   \_____________________/     \________________________/
#    Distribution DSL       RandomVariable DSL            Distribution DSL

where the .rv switches to the RandomVariable DSL and the .dist switches back.

I realize this is a significant change from your original mixin idea, but at least it keeps your nice method chaining syntax. Would this limited design something you'd still like to work on?

[eb8680]

@ecotner - @fritzo's proposal seems like a good first step, we often start new pyro.contrib modules when experimenting with new API designs. I agree with you that an interface like this is a worthy design goal, and I certainly don't want to discourage you from contributing!

[ecotner]

Thanks guys, those are definitely valid concerns. I wasn't even aware of funsor or your plans to integrate it into pyro, so I'm not surprised my idea could potentially break some of the future design elements you had in mind.

I like the idea of easily switching between Distribution <--> RandomVariable; and the syntax you outlined seems pretty slick itself! (getting strong pandas vibes). Since RV.__add__ could basically call RV.add under the hood, you would be free to use either method chaining or operator overloading as you saw fit. And I had no assumptions that the mixin idea would even be feasible; my original post was simply to illustrate the concept and give an example that demonstrated the desired functionality without making modifications to the pyro source code.

@eb8680 I'm not discouraged at all; I don't think I've ever had a PR accepted on the first try 😉.

[ecotner]

So, I moved everything over to pyro.contrib and added the method chaining, but I'm getting a problem which I believe to be a circular import error that is potentially preventing me from adding the .rv property to Distribution. Basically, RandomVariable needs TransformedDistribution, which needs Distribution, which needs RandomVariable... etc.

I found a way around it by moving the import inside the .rv function, but it seems a little hacky. Thoughts?

Also, do you guys follow the isort format anymore? I ran make format and it modified the import statements of like 130 different files. I'm pretty sure that's why my build is failing, but I didn't want to change so many files at once...

[fritzo]

This looks great! I think this only needs documentation. I've added a list of documentation tasks to your PR description.

[local import] seems a little hacky.

I'm ok with some hackiness connecting contrib code with non-contrib code. If you do find a better practice let me know.

do you guys follow the isort format anymore?

We don't enforce isort automatically, and I don't know an easy way to enforce subtler linting/formating requirements for community PRs. Also there are a couple files that I believe need non-isort import order, e.g. pyro/distributions/init.py and I haven't figured out how to annotate single lines or single files with isort. Again, if you know how to do this we'd appreciate any code cleanliness tips.

[fritzo]

@eb8680 are you ok with this experimental .rv interface? If so I'm ready to merge.

[eb8680]

Looks good to me