fix typo input->probs

synnada-ai · Feb 1, 2025 · b375b40 · b375b40
1 parent b3dae6a
commit b375b40
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mithril/backends/with_autograd/jax_backend/backend.py b/mithril/backends/with_autograd/jax_backend/backend.py
@@ -532,11 +532,11 @@ def multinomial(
 
             # 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