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

mithril/backends/with_autograd/jax_backend/backend.py

@@ -513,29 +513,30 @@
             replacement: whether to sample with replacement
         """
         prng_key = self._get_prng_key(key)
-        input = jax.numpy.asarray(probs)
-        input = input / jax.numpy.sum(input, axis=-1, keepdims=True)
-        batch_size = input.shape[:-1]
-        logits = jax.numpy.log(jax.numpy.maximum(input, 1e-37))
+        probs = jnp.asarray(probs)
+        probs = probs / jnp.sum(probs, axis=-1, keepdims=True)
+
+        # Mask zero probabilities to avoid log(0) without adding small constants
+        logits = jnp.where(probs > 0, jnp.log(probs), -jnp.inf)
 
         if replacement:
             # Use categorical directly - much faster than choice
             samples = jax.random.categorical(
                 prng_key,
                 logits,  # avoid log(0)
-                shape=batch_size + (num_samples,),
+                shape=probs.shape[:-1] + (num_samples,),
             )
         else:
             # TODO: This algorithm is not efficient for small num_samples
             # consider more efficient algorithm
 
             # For without replacement, use Gumbel-max trick
             # This is much faster than using choice
-            z = jax.random.gumbel(prng_key, shape=input.shape + (num_samples,))
+            z = jax.random.gumbel(prng_key, shape=probs.shape + (num_samples,))
             # Add log probabilities for Gumbel-max trick,
             z = z + logits[..., None]
             # Get top k indices
-            samples = jax.numpy.argsort(-z, axis=input.ndim - 1)[..., :num_samples, 0]
+            samples = jax.numpy.argsort(-z, axis=probs.ndim - 1)[..., :num_samples, 0]
 
         return samples