Feature request: Conditional sampling

In many TFP bayesian use cases, it's very helpful to specify a joint distribution using a JointDistribution* object - it makes sampling (for prior-predictive checks) straightforward (sorry), and exposes a log_prob function necessary for the sampler. However, since many (most?) of these cases involve some sort of conditioning, we end up writing a function closure which is very confusing and possibly error prone (mostly shape errors, but also type errors):

lp = lambda *x: model.log_prob(list(x) + [tf.cast(df['y'],tf.float32)[tf.newaxis, ...]])

Exposing some sort of conditioning method, instead, could be amazing.

For example, for a JointDistributionSequential (which is represented by a list), perhaps something along these lines:

posterior_lp = model.cond_lp(conditioned_var_idx = ..., condition_on = ...)

?

Thanks in advance!

Start by reviewing the existing JointDistribution* APIs and the closure-based log_prob example in the issue. Define the conditioning interface and its expected behavior for JointDistributionSequential, including how conditioned variables and values are represented; done means a supported conditional log_prob path replaces the error-prone closure pattern.