fix: prevent zero-probability events from being sampled in multinomial function in Jax-backend

[SinanGncgl]

Description

This PR fixes the issue #7 where zero-probability events could still be sampled due to the use of a small constant (1e-37) in the logits calculation. The fix replaces this constant with -inf to ensure such events are never selected and eliminate zero prob event.

What is Changed

Include the changes introduced in this PR.

• Updated the multinomial implementation to handle zero-probability outcomes correctly.
• Replaced 1e-37 with -inf using jnp.where for stable log calculations.

Checklist:

• Tests that cover the code added.
• Corresponding changes documented.
• All tests passed.
• The code linted and styled (pre-commit run --all-files has passed).

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 88.88%. Comparing base (5a07669) to head (3b14af7).
Report is 2 commits behind head on main.

@@            Coverage Diff             @@
##             main     #177      +/-   ##
==========================================
+ Coverage   88.71%   88.88%   +0.17%     
==========================================
  Files          61       61              
  Lines       15979    15968      -11     
==========================================
+ Hits        14175    14193      +18     
+ Misses       1804     1775      -29     

Files with missing lines	Coverage Δ
...hril/backends/with_autograd/jax_backend/backend.py	88.33% <100.00%> (+3.56%)	⬆️

... and 7 files with indirect coverage changes

[mehmetozsoy-synnada]

Instead of using jnp.where in order to avoid numerical instabilities, I tried this

logits = jax.numpy.log(probs)

I tested the edge cases in simple ways and did not encounter any problems as far as I observed. I think we could also write it this way, as it is simpler.

[ozankabak]

Seems like the log function automatically returns -inf when the input is zero.

[ozankabak]

Seems like we have a formatting problem, but it otherwise looks great. If you make sure the pre-commit procedure passes, we can proceed with the merge.

[ozankabak]

Thank you @SinanGncgl 🚀