Nans in JAX multinomial dispatch

[ricardoV94]

Description

PyMC tests are failing: pymc-devs/pymc#7740

Reproducible example:

import pytensor.tensor as pt

p = pt.eye(3)
rv = pt.random.multinomial(n=5, p=p)
rv.eval(mode="JAX")
# Array([[ 5., nan, nan],
#        [ 0.,  5.,  0.],
#        [ 0.,  0.,  5.]], dtype=float64)

I guess it could be a problem with binomial with p=0?

CC @educhesne

[educhesne]

I am very sorry about this... I think there is a division by zero happening here

[educhesne]

Do you know why it was not caught by the tests in the PR ?

[ricardoV94]

I guess we were not testing edge cases