Fix nan in jax implementation of Multinomial

Author: ricardoV94

pytensor/link/jax/dispatch/random.py

@@ -411,7 +411,7 @@ def sample_fn(rng_key, size, dtype, n, p):
         def _binomial_sample_fn(carry, p_rng):
             s, rho = carry
             p, rng = p_rng
-            samples = jax.random.binomial(rng, s, p / rho)
+            samples = jnp.where(p == 0, 0, jax.random.binomial(rng, s, p / rho))
             s = s - samples
             rho = rho - p
             return ((s, rho), samples)

tests/link/jax/test_random.py

@@ -733,6 +733,12 @@ def test_multinomial():
         samples.std(axis=0), np.sqrt(n[0, :, None] * p * (1 - p)), rtol=0.1
     )
 
+    # Test with p=0
+    g = pt.random.multinomial(n=5, p=pt.eye(4))
+    g_fn = compile_random_function([], g, mode="JAX")
+    samples = g_fn()
+    np.testing.assert_array_equal(samples, np.eye(4) * 5)
+
 
 @pytest.mark.skipif(not numpyro_available, reason="VonMises dispatch requires numpyro")
 def test_vonmises_mu_outside_circle():