Numpy fix for Multinomial sampling.

PiperOrigin-RevId: 346338660

Author: brianwa84

tensorflow_probability/python/distributions/multinomial.py

@@ -413,7 +413,7 @@ def fn(i, num_trials, consumed_prob, accum):
 
     num_trials = tf.cast(num_trials, probs.dtype)
     # Pre-broadcast with probs
-    num_trials += tf.zeros_like(probs[..., 0])
+    num_trials = num_trials + tf.zeros_like(probs[..., 0])
     # Pre-enlarge for different output samples
     num_trials = _replicate_along_left(num_trials, num_samples)
     i = tf.constant(0)