Tensorflow Probability: How can I calculate log_prob with weights for each feature of each sample?

Question

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 weight tensor of floats whose shape matches x i.e weights.shape=(batch_size, features). So this weight 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.

Any help is greatly appreciated!