The InvalidArgumentError occurred because dists(shape[10000]) could not broadcast directly with heights(shape[352]). This can be fixed by reshaping mu and sigma to [10000, 1], which enable implicit broadcasting to compute probabilities for all 352 heights against each of the 10,000 distributions and yields the desired [10000, 352] result without needing tf.map_fn or tf.vectorized_map. Please refer this gist for the mentioned approach.