Add a TFP JointDistribution instance for Edward2 models.

[edward-bot]

Add a TFP JointDistribution instance for Edward2 models.

[dustinvtran]

Thanks @davmre! It looks like unit tests are breaking. Does tfp.distributions.AutovectorizedJointDistribution not exist?

[davmre]

Oh gosh, I did not expect this to become a PR! Sorry for the confusion. I've been integrating TFP JointDistributions (tl;dr: these are just probabilistic models that fit into the Distributions API; the idea is that they provide a common interface for connecting model representations and inference routines) with tf.vectorized_map to try to make it cleaner to write models that support batched log-prob and sampling, in order to do multi-chain HMC, multi-sample VI bounds, etc. I prototyped an E2 version just to see how it'd work out, and it looks like creating an internal CL automatically creates a github PR?

Anyway thanks for the comments! This is definitely not ready for submission: AutovectorizedJointDistribution is indeed not checked into TFP yet, and the tests here are directly ripped off from JointDistributionCoroutine and probably need some revision. I expect AutovectorizedJointDistribution or something equivalent to hit TFP in the next week or two, so I'll follow up once this is actually runnable. :-)

For what it's worth, I'd say the main thing you currently get from using the JointDistribution API (aside from the subjective aesthetic satisfaction of treating models as 'just another distribution') is interoperability with TFP inference routines. For example, we can do chi-squared VI in one line with

tfp.vi.fit_surrogate_posterior(target_log_prob_fn=model.log_prob, 
                               surrogate_posterior=variational_model,
                               optimizer=tf.optimizers.Adam(0.1),
                               num_steps=200,
                               variational_loss_fn=functools.partial(
                                  tfp.vi.monte_carlo_variational_loss,
                                  discrepancy_fn=tfp.vi.chi_square))
posterior_samples = variational_model.sample(50)

by representing the model and variational_model as JointDistribution instances. In the near future JointDistribution will also do fun stuff like automatically vectorize your model with tf.vectorized_map (of course E2 could in principle do this itself in make_log_joint_fn) and apply default support bijectors.

[dustinvtran]

I prototyped an E2 version just to see how it'd work out, and it looks like creating an internal CL automatically creates a github PR?

If you run an internal presubmit (and Copybara passes), a PR will be created and/or updated.

Anyway thanks for the comments! This is definitely not ready for submission: AutovectorizedJointDistribution is indeed not checked into TFP yet, and the tests here are directly ripped off from JointDistributionCoroutine and probably need some revision. I expect AutovectorizedJointDistribution or something equivalent to hit TFP in the next week or two, so I'll follow up once this is actually runnable. :-)

Got it!

For what it's worth, I'd say the main thing you currently get from using the JointDistribution API (aside from the subjective aesthetic satisfaction of treating models as 'just another distribution') is interoperability with TFP inference routines. [...] In the near future JointDistribution will also do fun stuff like automatically vectorize your model with tf.vectorized_map (of course E2 could in principle do this itself in make_log_joint_fn) and apply default support bijectors.

Those benefits do sound like what you can already do given a log prob from make_log_joint_fn. Regarding subjective aesthetic satisfaction, distributions as objects is a natural idea for core building blocks like normal, chi2, etc. However, it seems like a lot of what one might expect starts to break down once you handle more complicated probability distributions such as conditionals, compound, autoregressive, and joints. As an abstraction, it adds a lot of magic that can make it feel unwieldy or inflexible in practice (that thunk for example, and the order it returns samples and how to manipulate the joint to do conditioning). I'm definitely open to learning more though!