Broadcast arrays manually in categorical sampling.

google · Jul 7, 2021 · f945982 · f945982
1 parent 56087dc
commit f945982
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/jax/_src/random.py b/jax/_src/random.py
@@ -1306,7 +1306,10 @@ def categorical(key: jnp.ndarray,
     _check_shape("categorical", shape, batch_shape)
 
   sample_shape = shape[:len(shape)-len(batch_shape)]
-  return jnp.argmax(gumbel(key, sample_shape + logits.shape, logits.dtype) + logits, axis=axis)
+  return jnp.argmax(
+      gumbel(key, sample_shape + logits.shape, logits.dtype) +
+      lax.expand_dims(logits, tuple(range(len(sample_shape)))),
+      axis=axis)
 
 
 def laplace(key: jnp.ndarray,

diff --git a/tests/random_test.py b/tests/random_test.py
@@ -379,6 +379,7 @@ def testBernoulli(self, p, dtype):
     ]
     for sample_shape in [(10000,), (5000, 2)]
     for dtype in jtu.dtypes.floating))
+  @jtu.disable_implicit_rank_promotion
   def testCategorical(self, p, axis, dtype, sample_shape):
     key = random.PRNGKey(0)
     p = np.array(p, dtype=dtype)