Fix nan for valid parameters in jax implementation of Multinomial

[ricardoV94]

Closes #1327

---

📚 Documentation preview 📚: https://pytensor--1328.org.readthedocs.build/en/1328/

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 82.01%. Comparing base (0b56ed9) to head (bd85204).
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1328   +/-   ##
=======================================
  Coverage   82.01%   82.01%           
=======================================
  Files         203      203           
  Lines       48805    48805           
  Branches     8688     8688           
=======================================
  Hits        40026    40026           
  Misses       6627     6627           
  Partials     2152     2152           

Files with missing lines	Coverage Δ
pytensor/link/jax/dispatch/random.py	94.01% <100.00%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[educhesne]

Your fix solves the problem when p is valid (the sum of p is 1), but if not you may have rho==0 without p==0; otherwise it is possible to do:
samples = jax.random.binomial(rng, s, jnp.where(rho == 0, 0, p / rho))

[ricardoV94]

Ah the problem is the division by zero, yes that's better. I guess I should also test n=0 multinomial

[ricardoV94]

Your fix solves the problem when p is valid (the sum of p is 1), but if not you may have rho==0 without p==0; otherwise it is possible to do: samples = jax.random.binomial(rng, s, jnp.where(rho == 0, 0, p / rho))

Actually I don't think I want to mask that, if p adds to more than 1, then it's fine to get nans. In numpy it would raise an error, but jax doesn't have runtime errors.