Return ed.RandomVariables in initializers and operate on them in regu…

[dustinvtran]

…larizers.

@markvdw @dusenberrymw @danijar @balajiln: PTAL. After approval, I will merge the PR via an internal process.

@afrozenator: I'm sending this as a GitHub pull request for external collaborators to review and add comments. I'm not sure this change works with Travis, as the test now requires tfp as an import. Has that Q been resolved?

[dusenberrymw]

This LGTM. Using Edward2 RVs makes this much cleaner, particularly for the KL divergences.

[dusenberrymw]

I was thinking that we could create and store this in the build function to remove a bit of Python overhead in eager mode, but it looks like Edward2 RVs only sample from the underlying distribution on the first access in eager mode due to rv.value using a memoization strategy by storing the sample on the first access. It seems like it could potentially simplify some use cases if rv.value called rv.distribution.sample on each access to rv.value. That being said, you definitely thought much longer on the design of Edward[1|2], so I bet there is a reason that would make this undesirable.

[dustinvtran]

We memoize for usability reasons, so that you operate on RandomVariables as if they were any other Tensor. For example, if you call x = ed.Normal(...), you would expect y = x + 1 and z = x + 1 to return the same result.

On caching: There's an issue that actually arises before Ed2, which is that during __init__, Distributions call TF ops on its parameters. This breaks gradients if Distribution.__init__ isn't recorded in the gradient tape. @vdumoulin raised a point about this in Brain Chat.

[danijar]

Any reason to not create this object in __init__ already instead of saving self.mean and save.stddev separately?

[dustinvtran]

Breaks Eager unfortunately if self.mean and self.stddev depend on trainable variables. See the above comment about caching. It would be nice if Distribution __init__s don't add TF ops to the graph.