MixtureSameFamily and BatchReshape don't mix well

Python 3.6.5
tensorflow==2.4.0
tensorflow-probability==0.12.1

I'm trying to create a mixture distribution with MixtureSameFamily after having reshaped the batch of the components distribution with BatchReshape. But calling log_prob on the resulting mixture distribution raises a NotImplementedError. I was just wondering if that functionality is almost available or far down the road.

Here's a minimal example of what I'm trying to do:

n_rows = 3
n_cols = 2
event_shape = 4
prediction = tfd.Independent(tfd.Normal((np.arange(event_shape).astype(np.float32) + 0.1), np.ones((n_rows, n_cols, 1), dtype=np.float32)), reinterpreted_batch_ndims=1)
prediction
==> <tfp.distributions.Independent 'IndependentNormal' batch_shape=[3, 2] event_shape=[4] dtype=float32>

prediction = tfp.distributions.BatchReshape(
    prediction, (n_rows * n_cols,), validate_args=True
)
prediction
==> <tfp.distributions.BatchReshape 'BatchReshapeIndependentNormal' batch_shape=[6] event_shape=[4] dtype=float32>

prediction.log_prob(np.ones((6, event_shape)))  # Works

cat = tfp.distributions.Categorical(
    probs=tf.ones((n_rows * n_cols,))
    / (n_rows * n_cols)
)
predictive_mixture = tfp.distributions.MixtureSameFamily(cat, prediction)
predictive_mixture
==> <tfp.distributions.MixtureSameFamily 'MixtureSameFamily' batch_shape=[] event_shape=[4] dtype=float32>

sample = predictive_mixture.sample(seed=0) # Sampling works
sample
==> <tf.Tensor: shape=(4,), dtype=float32, numpy=array([0.42147806, 0.25534433, 2.05793   , 2.4141035 ], dtype=float32)>

predictive_mixture.log_prob(sample)
==> NotImplementedError: Broadcasting is not supported; unexpected batch and event shape (expected [6 4], saw [1 4])

Start by running the minimal MixtureSameFamily and BatchReshape example and inspect MixtureSameFamily.log_prob together with BatchReshape. Done means the reproduced mixture accepts the reshaped component batch and log_prob returns a result instead of raising the reported NotImplementedError.