More stable fix for JAX Multinomial

[ricardoV94]

#1328 was not good enough for the test in PyMC :(

I changed it to stop taking draws once remaining_n drops to zero. This should be more robust against numerical precision issues for valid parameters. It's also what our numba implementation is doing, so at least they are aligned:

pytensor/pytensor/link/numba/dispatch/random.py

Lines 166 to 167 in 75789de

	if remaining_n <= 0:
	break

Or I should perhaps just gone with the suggestion @educhesne gave.

I confirm PyMC is now happy

---

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

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 82.01%. Comparing base (afb7695) to head (3debfdc).
Report is 2 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1339   +/-   ##
=======================================
  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% <ø> (ø)

🚀 New features to boost your workflow:

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

[ricardoV94]

I'm gonna go away and merge this, it has slowed down things enough as it is :)