More stable fix for JAX Multinomial

Author: ricardoV94

Diff for: pytensor/link/jax/dispatch/random.py

@@ -412,7 +412,9 @@ def _binomial_sample_fn(carry, p_rng):
             remaining_n, remaining_p = carry
             p, rng = p_rng
             samples = jnp.where(
-                p == 0, 0, jax.random.binomial(rng, remaining_n, p / remaining_p)
+                remaining_n == 0,
+                0,
+                jax.random.binomial(rng, remaining_n, p / remaining_p),
             )
             remaining_n -= samples
             remaining_p -= p