I am working on TF Probability tutorials and working on a 2D grid approximation for mu, sigma of Normal distribution. I am trying to understand to what level I should expect broadcasting to work for the following code:

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions

mu_list = tf.linspace(start=150, stop=160, num=100)
sigma_list = tf.linspace(start=7, stop=9, num=100)
mesh = tf.meshgrid(mu_list, sigma_list)
mu = tf.cast(tf.reshape(mesh[0], -1), tf.float32)
sig = tf.cast(tf.reshape(mesh[1], -1), tf.float32)
dists = tfd.Normal(loc=mu, scale=sig)
heights = tfd.Normal(loc=150, scale=20).sample(352)
dists.prob(heights)

This fails with

Incompatible shapes: [352] vs. [10000]

I am pretty sure I can solve this with either a tf.map_fn or a tf.vectorized_map operation, but curious if it is possible to generate a [10000, 352] shape tensor in the .prob/.log_prob call.