Feature request: 2-D weights option for log prob

I am working on a VAE with a Gaussian decoder using tensorflow probability. I calculate the log likelihood using x_hat.log_prob(x) where x_hat is the output of the decoder (tfp.distributions.Distribution). However, in my application it is important to apply weights to the values from the log_prob summation using a weights tensor of floats whose shape matches x_true i.e weights.shape=(batch_size, features). So this weights tensor will apply a different weight to each feature of each sample.

Essentially I want to do what is being done in the MSE example below:

MSE = tf.reduce_sum(0.5 * weights * (x - recon_x)**2)

where x, recon_x and weights all have shape (batch_size, features).

I tried looking for a way to get the log_prob before the summation so that I can apply the weights and then sum but I couldn’t find something relevant in the documentation.

Is there another way to go about doing this?

Start by reading the TensorFlow Probability Distribution.log_prob documentation and the behavior of x_hat.log_prob(x) for batched feature inputs. Determine how per-feature log probabilities and a weights tensor shaped like x_true could be supported, then verify that weighted values can be summed per sample as requested.