Issue when calling log_prob attribute in a logistic mixture distribution instance

Hey,

I have the following issue when using this particular instance of mixture same family distribution. Here is a snipper of my code:

`mean, scale, logits = self.decode(z)

  discretized_logistic_dist = tfd.QuantizedDistribution(
    distribution=tfd.TransformedDistribution(
        distribution=tfd.Logistic(loc=mean, scale=scale),
        bijector=tfb.Shift(shift=-0.5)),
    low=0.,
    high=2**16 - 1.)

  logistic_dist = tfd.Logistic(loc=mean, scale=scale)


  mixture_dist = tfd.MixtureSameFamily(
  mixture_distribution=tfd.Categorical(logits=logits),
  components_distribution=logistic_dist)

  ll = -tf.reduce_sum(mixture_dist.log_prob(x), axis=[1,2,3])`

I am basically using a VAE to learn the parameters of a logistic mixture family. The issue if the following:

when calling mixture_dist.log_prob(x) on my input data, in this case cifar10, for some odd reason the input is being reshaped into (x.shape, 1) basically performing tf.expand_dims on x at least assuming that is the case. Thus, an error of invalid dimensions is thrown at me

Dimensions must be equal, but are 32 and 256 for '{{node MixtureSameFamily_1/log_prob/Logistic_1/log_prob/sub}} = Sub[T=DT_FLOAT](MixtureSameFamily_1/log_prob/pad_sample_dims/Reshape, split_1)' with input shapes: [256,32,32,3,1], [256,32,32,15].

I have tried several hacks using tf.squeeze and converting the log_prob instance into a function prior to applying x in i, but to no result.

Could you please assist with the issue.

The report provides no repository file or test. Start by reproducing the shown TensorFlow Probability MixtureSameFamily.log_prob call with CIFAR10-shaped input and inspect the reported pad_sample_dims/Reshape operation. Done means the dimensionality behavior is isolated and covered by an appropriate regression test or clarified usage guidance.