InvalidArgumentError when evaluating logp on inits sampled from model

In previous versions of TFP, I was able to sample from my model to get inits for MCMC sampling, so I could evaluate the log_prob directly:

inits = model.sample(n_chains)
model.log_prob(inits)

However, in '0.10.0-dev20200413' (TF version '2.2.0-dev20200413'), the call to log_prob fails with the following exception:

InvalidArgumentError: Index out of range using input dim 0; input has only 0 dims [Op:StridedSlice] name: strided_slice/

I assume this has something to do with the shape of the inputs? These are samples from a JointDistributionSequential and the shapes on the sampled values are:

[TensorShape([4]),
 TensorShape([4, 50]),
 TensorShape([4]),
 TensorShape([4, 3]),
 TensorShape([4]),
 TensorShape([4, 3]),
 TensorShape([4]),
 TensorShape([4, 7]),
 TensorShape([4]),
 TensorShape([4, 13]),
 TensorShape([4, 1200])]

Start by reproducing the reported JointDistributionSequential sample/log_prob sequence with TensorFlow 2.2.0-dev20200413 and TensorFlow Probability 0.10.0-dev20200413. Inspect the listed sample shapes and the failing log_prob path; done means model.log_prob(model.sample(n_chains)) evaluates without the reported InvalidArgumentError.